fixed inference bug

inference.py

@@ -3,6 +3,7 @@ import joblib
 import os
 import numpy as np
 
+
 class CameraMotionCompensator:
     def __init__(self):
         self.prev_frame = None
@@ -41,6 +42,7 @@ class CameraMotionCompensator:
         self.prev_desc = desc
         return M
 
+
 class ImprovedSlidingWindowTracker:
     def __init__(self, scale_factor=2.0, overlap=0.3):
         self.scale_factor = scale_factor
@@ -73,26 +75,35 @@ class ImprovedSlidingWindowTracker:
         return windows
 
     def score_window(self, gray, window, template, template_desc):
-        x,y,w,h = map(int,window)
-        roi = gray[y:y+h,x:x+w]
-        if roi.shape[0]<20 or roi.shape[1]<20:
-            return 0
-        roi = cv2.resize(roi,(template.shape[1],template.shape[0]))
-        _,desc = self.sift.detectAndCompute(roi,None)
-        if desc is None or template_desc is None:
-            return 0
-        matches = self.flann.knnMatch(template_desc,desc,k=2)
-        good = [m for m,n in matches if m.distance < 0.7*n.distance]
-        if not good:
+        try:
+            x,y,w,h = map(int,window)
+            if y+h > gray.shape[0] or x+w > gray.shape[1]:
+                return 0
+            roi = gray[y:y+h,x:x+w]
+            if roi.shape[0]<20 or roi.shape[1]<20:
+                return 0
+            roi = cv2.resize(roi,(template.shape[1],template.shape[0]))
+            _,desc = self.sift.detectAndCompute(roi,None)
+            if desc is None or template_desc is None or len(desc)==0:
+                return 0
+            matches = self.flann.knnMatch(template_desc,desc,k=2)
+            good = [m for match_pair in matches if len(match_pair)==2 for m,n in [match_pair] if m.distance < 0.7*n.distance]
+            if not good:
+                return 0
+            return len(good)*(1-np.mean([m.distance for m in good])/512)
+        except:
             return 0
-        return len(good)*(1-np.mean([m.distance for m in good])/512)
+
 
 class ObjectTrackerInference:
     def __init__(self, model_dir):
+        print(f"Loading models from {model_dir}...")
         self.position_model = joblib.load(os.path.join(model_dir,'position_model.joblib'))
         self.size_model = joblib.load(os.path.join(model_dir,'size_model.joblib'))
         self.position_scaler = joblib.load(os.path.join(model_dir,'position_scaler.joblib'))
         self.size_scaler = joblib.load(os.path.join(model_dir,'size_scaler.joblib'))
+        print("Models loaded successfully!")
+        
         self.window_tracker = ImprovedSlidingWindowTracker()
         self.motion = CameraMotionCompensator()
         self.template = None
@@ -122,21 +133,24 @@ class ObjectTrackerInference:
 
         if self.template is None:
             x,y,w,h = map(int,prev_bbox)
-            self.template = gray[y:y+h,x:x+w]
+            x = max(0, min(x, gray.shape[1]-w))
+            y = max(0, min(y, gray.shape[0]-h))
+            self.template = gray[y:y+h,x:x+w].copy()
             _,self.template_desc = self.window_tracker.sift.detectAndCompute(self.template,None)
 
         best_score = -1
-        best_window = None
+        best_window = prev_bbox
         for w in windows:
             s = self.window_tracker.score_window(gray,w,self.template,self.template_desc)
             if s > best_score:
                 best_score = s
                 best_window = w
 
-        if best_window is None:
-            x,y,w,h = map(int,prev_bbox)
-        else:
-            x,y,w,h = map(int,best_window)
+        x,y,w,h = map(int,best_window)
+        x = max(0, min(x, gray.shape[1]-10))
+        y = max(0, min(y, gray.shape[0]-10))
+        w = min(w, gray.shape[1]-x)
+        h = min(h, gray.shape[0]-y)
 
         roi = cv2.resize(gray[y:y+h,x:x+w],(64,64))
         hog = cv2.HOGDescriptor((64,64),(16,16),(8,8),(8,8),9).compute(roi).flatten()[:64]
@@ -161,60 +175,99 @@ class ObjectTrackerInference:
         inter=(xr-xl)*(yb-yt)
         return inter/(w1*h1+w2*h2-inter)
 
-    def track_video(self, video_path, init_bbox, output):
-        cap=cv2.VideoCapture(video_path)
-        w,h=int(cap.get(3)),int(cap.get(4))
-        out=cv2.VideoWriter(output,cv2.VideoWriter_fourcc(*'mp4v'),30,(w,h))
-        cur=init_bbox
-        frame_idx=0
-
-        while True:
-            ret,frame=cap.read()
-            if not ret:
-                break
-
-            M=self.motion.estimate_motion(frame)
-            feats,search_bbox,windows=self.extract_features(frame,cur,M)
-
-            pos=self.position_model.predict(self.position_scaler.transform(feats))
-            size=self.size_model.predict(self.size_scaler.transform(feats))
-            pred=[int(pos[0,0]),int(pos[0,1]),int(size[0,0]),int(size[0,1])]
-
-            self.template_update_counter+=1
-            if self.template_update_counter>=5 and self.prev_bbox is not None:
-                if self.calculate_iou(self.prev_bbox,pred)>0.6:
-                    g=cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
-                    x,y,w1,h1=pred
-                    self.template=g[y:y+h1,x:x+w1]
-                    _,self.template_desc=self.window_tracker.sift.detectAndCompute(self.template,None)
-                    self.template_update_counter=0
-
-            for wx,wy,ww,wh in windows:
-                cv2.rectangle(frame,(wx,wy),(wx+ww,wy+wh),(0,255,255),1)
-
-            hh,ww=frame.shape[:2]
-            for yy in range(0,hh,32):
-                for xx in range(0,ww,32):
-                    sp=np.array([xx,yy,1])
-                    ep=np.dot(M,sp)
-                    if abs(ep[0]-xx)>1 or abs(ep[1]-yy)>1:
-                        cv2.arrowedLine(frame,(xx,yy),(int(ep[0]),int(ep[1])),(0,255,0),1,tipLength=0.2)
-
-            x,y,w1,h1=pred
-            cv2.rectangle(frame,(x,y),(x+w1,y+h1),(0,255,0),2)
-            cv2.putText(frame,f'Frame: {frame_idx}',(10,30),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),2)
-
-            out.write(frame)
-            self.prev_bbox=pred
-            cur=pred
-            frame_idx+=1
-
-        cap.release()
-        out.release()
+    def track_video(self, video_path, init_bbox, output_path='tracked_output.mp4'):
+        print(f"Opening video: {video_path}")
+        
+        try:
+            cap=cv2.VideoCapture(video_path)
+            if not cap.isOpened():
+                raise ValueError(f"Cannot open video: {video_path}")
+            
+            w,h=int(cap.get(3)),int(cap.get(4))
+            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+            print(f"Video: {w}x{h}, {total_frames} frames")
+            
+            fourcc = cv2.VideoWriter_fourcc(*'avc1')
+            out=cv2.VideoWriter(output_path, fourcc, 30, (w,h))
+            
+            if not out.isOpened():
+                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+                out=cv2.VideoWriter(output_path, fourcc, 30, (w,h))
+            
+            # Reset state
+            self.motion.prev_frame = None
+            self.template = None
+            self.template_desc = None
+            self.prev_bbox = None
+            self.template_update_counter = 0
+            
+            cur=init_bbox
+            frame_idx=0
+
+            print("Starting tracking...")
+            while True:
+                ret,frame=cap.read()
+                if not ret:
+                    break
+
+                M=self.motion.estimate_motion(frame)
+                feats,search_bbox,windows=self.extract_features(frame,cur,M)
+
+                pos=self.position_model.predict(self.position_scaler.transform(feats))
+                size=self.size_model.predict(self.size_scaler.transform(feats))
+                pred=[int(pos[0,0]),int(pos[0,1]),int(size[0,0]),int(size[0,1])]
+
+                self.template_update_counter+=1
+                if self.template_update_counter>=5 and self.prev_bbox is not None:
+                    if self.calculate_iou(self.prev_bbox,pred)>0.6:
+                        g=cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
+                        x,y,w1,h1=pred
+                        self.template=g[y:y+h1,x:x+w1].copy()
+                        _,self.template_desc=self.window_tracker.sift.detectAndCompute(self.template,None)
+                        self.template_update_counter=0
+
+                # Draw yellow search windows
+                for wx,wy,ww,wh in windows:
+                    cv2.rectangle(frame,(wx,wy),(wx+ww,wy+wh),(0,255,255),1)
+
+                # Draw green motion arrows
+                hh,ww=frame.shape[:2]
+                for yy in range(0,hh,32):
+                    for xx in range(0,ww,32):
+                        sp=np.array([xx,yy,1])
+                        ep=np.dot(M,sp)
+                        if abs(ep[0]-xx)>1 or abs(ep[1]-yy)>1:
+                            cv2.arrowedLine(frame,(xx,yy),(int(ep[0]),int(ep[1])),(0,255,0),1,tipLength=0.2)
+
+                # Draw tracked bounding box
+                x,y,w1,h1=pred
+                cv2.rectangle(frame,(x,y),(x+w1,y+h1),(0,255,0),2)
+                cv2.putText(frame,f'Frame: {frame_idx}',(10,30),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,255),2)
+
+                out.write(frame)
+                self.prev_bbox=pred
+                cur=pred
+                frame_idx+=1
+                
+                if frame_idx % 30 == 0:
+                    print(f"Processed {frame_idx}/{total_frames} frames")
+
+            cap.release()
+            out.release()
+            
+            print(f"Tracking complete! Saved to: {output_path}")
+            return output_path
+            
+        except Exception as e:
+            print(f"Error during tracking: {str(e)}")
+            raise
+
 
 def main():
     tracker=ObjectTrackerInference('models')
-    tracker.track_video('input_video.mp4',[100,100,50,50],'tracked_output.mp4')
+    result = tracker.track_video('input_video.mp4',[100,100,50,50],'tracked_output.mp4')
+    print(f"Output: {result}")
+
 
 if __name__=="__main__":
     main()