YOLODetector.java

package yolo;

import org.bytedeco.opencv.opencv_core.*;
import static org.bytedeco.opencv.global.opencv_imgcodecs.imwrite;
import static org.bytedeco.opencv.global.opencv_imgcodecs.imread;
import static org.bytedeco.opencv.global.opencv_imgproc.rectangle;
import static org.bytedeco.opencv.global.opencv_imgproc.LINE_8;

import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.io.File;
import java.io.IOException;
import java.io.FileWriter;

public class YOLODetector {
    private static YOLONet yoloNet;

    public static void main(String[] args) {
        String keyFramesDir = "src/main/java/vslam/KeyFrames";
        String outputDir = "src/main/java/vslam/BoundedKeyFrames";
        String csvDir = "src/main/java/vslam/BoundedInfo";
        Path outputPath = Paths.get(outputDir);
        Path csvPath = Paths.get(csvDir);

        // Ensure output directories exist
        try {
            if (!Files.exists(outputPath)) {
                Files.createDirectories(outputPath);
            }
            if (!Files.exists(csvPath)) {
                Files.createDirectories(csvPath);
            }
        } catch (Exception e) {
            e.printStackTrace();
            return;
        }

        yoloNet = new YOLONet(
                "src/main/java/yolo/yolov4.cfg",
                "src/main/java/yolo/yolov4.weights",
                "src/main/java/yolo/coco.names",
                608, 608); // Adjust size if necessary for performance

        if (!yoloNet.setup()) {
            System.err.println("Failed to setup YOLONet");
            return;
        }

        try (Stream<Path> paths = Files.walk(Paths.get(keyFramesDir))) {
            List<String> files = paths.filter(Files::isRegularFile)
                                      .map(Path::toString)
                                      .collect(Collectors.toList());

            for (String file : files) {
                Mat image = imread(file);
                List<YOLONet.ObjectDetectionResult> results = yoloNet.predict(image);
                File csvFile = new File(csvPath + "/" + new File(file).getName().replace(".png", ".csv")); // Assuming PNG images
                try (FileWriter writer = new FileWriter(csvFile)) {
                    writer.write("Class,Confidence,X,Y,Width,Height\n");

                    for (YOLONet.ObjectDetectionResult result : results) {
                        rectangle(image, new Point(result.x, result.y),
                                  new Point(result.x + result.width, result.y + result.height),
                                  Scalar.MAGENTA, 2, LINE_8, 0);
                        writer.write(String.format("%s,%f,%d,%d,%d,%d\n", 
                            result.className, result.confidence, result.x, result.y, result.width, result.height));
                    }
                }

                String outputPathString = outputDir + "/" + new File(file).getName();
                imwrite(outputPathString, image);
                System.out.println("Processed and saved: " + outputPathString);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}