From c0ea23f0a9d7ba2b0bafb429818e517a6bc33237 Mon Sep 17 00:00:00 2001
From: kishor G <65946231+kishorsumathi@users.noreply.github.com>
Date: Wed, 4 Oct 2023 18:42:48 +0530
Subject: [PATCH] Update mineral.py

---
 src/margsoft/mineral.py | 41 ++++++++++++++++++++++-------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/src/margsoft/mineral.py b/src/margsoft/mineral.py
index 06f700d..7b367ed 100644
--- 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()