Update mineral.py

Codebugged-Research · Oct 4, 2023 · c0ea23f · c0ea23f
1 parent 3e06952
commit c0ea23f
Showing 1 changed file with 22 additions and 19 deletions.
diff --git a/src/margsoft/mineral.py b/src/margsoft/mineral.py
@@ -43,26 +43,29 @@ def collectdataset(rtsp,names_file,weight_file,cfg_file,path,dir_n,x_1,y_1,w_1,z
             print(f"CPU utilization: {cpu_percent}%")
             print(f"CPU utilization: {cpu_percent}%")
             ret, frame = cap.read()
-            frame=cv2.resize(frame, (900, 900))
-            if poly:
-                frame_crop=frame.copy()
-                cv2.fillPoly(frame_crop, [pts_left], 0)
-                cv2.fillPoly(frame_crop, [pts_right], 0)
-            elif roi:
-                frame_crop=frame[y_1:z_1,x_1:w_1]
+            if ret:
+                frame=cv2.resize(frame, (900, 900))
+                if poly:
+                    frame_crop=frame.copy()
+                    cv2.fillPoly(frame_crop, [pts_left], 0)
+                    cv2.fillPoly(frame_crop, [pts_right], 0)
+                elif roi:
+                    frame_crop=frame[y_1:z_1,x_1:w_1]
+                else:
+                    frame_crop=frame
+                frame_crop= np.array(frame_crop)
+                # frame=cv2.resize(frame, (900, 900))
+                frame = np.array(frame)
+                classes, scores, boxes = model.detect(frame_crop, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)
+                for (classid, score, box) in zip(classes, scores, boxes):
+                        print(class_names[classid],score)
+                        print("found!!!")
+                        curr_datetime = datetime.utcnow().strftime('%Y%m%d%H%M%S%f')
+                        f_name = path+dir_n+"/"+str(curr_datetime)+".jpg"
+                        cv2.imwrite(f_name, frame)
+                        print(f_name)
             else:
-                frame_crop=frame
-            frame_crop= np.array(frame_crop)
-            # frame=cv2.resize(frame, (900, 900))
-            frame = np.array(frame)
-            classes, scores, boxes = model.detect(frame_crop, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)
-            for (classid, score, box) in zip(classes, scores, boxes):
-                    print(class_names[classid],score)
-                    print("found!!!")
-                    curr_datetime = datetime.utcnow().strftime('%Y%m%d%H%M%S%f')
-                    f_name = path+dir_n+"/"+str(curr_datetime)+".jpg"
-                    cv2.imwrite(f_name, frame)
-                    print(f_name)
+                print("No Frame")
     except KeyboardInterrupt:
         print("Bye")
         sys.exit()