clc;
clearvars;
close all;

videoFile = 'dog_bump.mp4'; 
% videoFile = 'dog_walk.mp4'; 

video = VideoReader(videoFile);


% Load a pre-trained YOLOv4 object detector
model = yolov4ObjectDetector('tiny-yolov4-coco');

% Read a frame from the video
frame = readFrame(video);

% Detect objects
[bboxes, scores, labels] = detect(model, frame);

% Filter for "dog" class
dogIdx = labels == "dog";
dogBboxes = bboxes(dogIdx, :);
dogScores = scores(dogIdx);

% Show detection
annotatedFrame = insertObjectAnnotation(frame, 'rectangle', dogBboxes, dogScores);
imshow(annotatedFrame);


%% Full video

outputVideo = VideoWriter('tracked_puppy', 'MPEG-4'); 
outputVideo.FrameRate = video.FrameRate;

open(outputVideo); 

frame_num = 0;
while hasFrame(video)
    frame = readFrame(video); 
    % frame = imrotate(frame, -90); % Uncomment for dog_bump

    % Detect objects
    [bboxes, scores, labels] = detect(model, frame);
    
    % Filter for "dog" class
    dogIdx = labels == "dog";
    dogBboxes = bboxes(dogIdx, :);
    dogScores = scores(dogIdx);

    annotatedFrame = insertObjectAnnotation(frame, 'rectangle', dogBboxes, dogScores);

    writeVideo(outputVideo, annotatedFrame); 
    % disp(labels);
    disp(frame_num);
    frame_num = frame_num + 1;
end

close(outputVideo); 
disp('Processed video saved.');

