Update app.py

app.py

@@ -1,4 +1,142 @@
-import matplotlib.pyplot as plt
+# import matplotlib.pyplot as plt
+# import numpy as np
+# from six import BytesIO
+# from PIL import Image
+# import tensorflow as tf
+# from object_detection.utils import label_map_util
+# from object_detection.utils import visualization_utils as viz_utils
+# from object_detection.utils import ops as utils_op
+# import tarfile
+# import wget 
+# import gradio as gr
+# from huggingface_hub import snapshot_download
+# import os
+
+# # Install TensorFlow within the Hugging Face environment
+# os.system('pip install tensorflow')
+
+# # Now you can import TensorFlow
+# import tensorflow as tf
+
+
+# PATH_TO_LABELS = 'data/label_map.pbtxt'   
+# category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
+
+# def pil_image_as_numpy_array(pilimg):
+
+#     img_array = tf.keras.utils.img_to_array(pilimg)
+#     img_array = np.expand_dims(img_array, axis=0)
+#     return img_array
+    
+# def load_image_into_numpy_array(path):
+                                    
+#     image = None
+#     image_data = tf.io.gfile.GFile(path, 'rb').read()
+#     image = Image.open(BytesIO(image_data))
+#     return pil_image_as_numpy_array(image)            
+
+# def load_model():
+#     download_dir = snapshot_download(REPO_ID)
+#     saved_model_dir = os.path.join(download_dir, "saved_model")
+#     detection_model = tf.saved_model.load(saved_model_dir)
+#     return detection_model
+
+# def load_model2():
+#     wget.download("https://nyp-aicourse.s3-ap-southeast-1.amazonaws.com/pretrained-models/balloon_model.tar.gz")
+#     tarfile.open("balloon_model.tar.gz").extractall()
+#     model_dir = 'saved_model'    
+#     detection_model = tf.saved_model.load(str(model_dir))
+#     return detection_model    
+
+# # samples_folder = 'test_samples
+# # image_path = 'test_samples/sample_balloon.jpeg
+# # 
+
+# def predict(pilimg):
+
+#     image_np = pil_image_as_numpy_array(pilimg)
+#     return predict2(image_np)
+
+# def predict2(image_np):
+
+#     results = detection_model(image_np)
+
+#     # different object detection models have additional results
+#     result = {key:value.numpy() for key,value in results.items()}
+    
+#     label_id_offset = 0
+#     image_np_with_detections = image_np.copy()
+
+#     viz_utils.visualize_boxes_and_labels_on_image_array(
+#         image_np_with_detections[0],
+#         result['detection_boxes'][0],
+#         (result['detection_classes'][0] + label_id_offset).astype(int),
+#         result['detection_scores'][0],
+#         category_index,
+#         use_normalized_coordinates=True,
+#         max_boxes_to_draw=200,
+#         min_score_thresh=.60,
+#         agnostic_mode=False,
+#         line_thickness=2)
+
+#     result_pil_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
+    
+#     return result_pil_img
+
+# import cv2
+
+# def predict_video(video_path):
+#     cap = cv2.VideoCapture(video_path)
+#     frame_width = int(cap.get(3))
+#     frame_height = int(cap.get(4))
+
+#     # Define the codec and create a video writer object
+#     out = cv2.VideoWriter('output.avi', cv2.VideoWriter_fourcc('M','J','P','G'), 10, (frame_width, frame_height))
+
+#     while cap.isOpened():
+#         ret, frame = cap.read()
+#         if not ret:
+#             break
+
+#         # Convert the frame to PIL image
+#         pil_image = Image.fromarray(frame)
+        
+#         # Perform object detection on the frame
+#         result_pil_img = predict(pil_image)
+
+#         # Convert the result back to a NumPy array
+#         result_np_img = tf.keras.utils.img_to_array(result_pil_img)
+
+#         # Write the frame with detected objects to the video output
+#         out.write(result_np_img.astype('uint8'))
+
+#     # Release the video capture and writer objects
+#     cap.release()
+#     out.release()
+
+#     return "output.avi"
+
+
+
+# REPO_ID = "Louisw3399/burgerorfriesdetector"
+# detection_model = load_model()
+# # pil_image = Image.open(image_path)
+# # image_arr = pil_image_as_numpy_array(pil_image)
+
+# # predicted_img = predict(image_arr)
+# # predicted_img.save('predicted.jpg')
+
+# gr.Interface(fn=predict,
+#              inputs=gr.Image(type="pil"),
+#              outputs=gr.Image(type="pil")
+#              ).launch(share=True)
+
+# gr.Interface(
+#     fn=predict_video,
+#     inputs=gr.Video(type="file", label="Upload a video"),
+#     outputs=gr.Video(type="file", label="Download the processed video")
+# ).launch(share=True)
+
 import numpy as np
 from six import BytesIO
 from PIL import Image
@@ -6,11 +144,9 @@ import tensorflow as tf
 from object_detection.utils import label_map_util
 from object_detection.utils import visualization_utils as viz_utils
 from object_detection.utils import ops as utils_op
-import tarfile
-import wget 
-import gradio as gr
 from huggingface_hub import snapshot_download
-import os
+import gradio as gr
+import cv2
 
 # Install TensorFlow within the Hugging Face environment
 os.system('pip install tensorflow')
@@ -18,19 +154,15 @@ os.system('pip install tensorflow')
 # Now you can import TensorFlow
 import tensorflow as tf
 
-
 PATH_TO_LABELS = 'data/label_map.pbtxt'   
 category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
 
 def pil_image_as_numpy_array(pilimg):
-
     img_array = tf.keras.utils.img_to_array(pilimg)
     img_array = np.expand_dims(img_array, axis=0)
     return img_array
     
 def load_image_into_numpy_array(path):
-                                    
-    image = None
     image_data = tf.io.gfile.GFile(path, 'rb').read()
     image = Image.open(BytesIO(image_data))
     return pil_image_as_numpy_array(image)            
@@ -41,28 +173,15 @@ def load_model():
     detection_model = tf.saved_model.load(saved_model_dir)
     return detection_model
 
-def load_model2():
-    wget.download("https://nyp-aicourse.s3-ap-southeast-1.amazonaws.com/pretrained-models/balloon_model.tar.gz")
-    tarfile.open("balloon_model.tar.gz").extractall()
-    model_dir = 'saved_model'    
-    detection_model = tf.saved_model.load(str(model_dir))
-    return detection_model    
-
-# samples_folder = 'test_samples
-# image_path = 'test_samples/sample_balloon.jpeg
-# 
-
 def predict(pilimg):
-
     image_np = pil_image_as_numpy_array(pilimg)
-    return predict2(image_np)
-
-def predict2(image_np):
+    return predict_objects(image_np)
 
+def predict_objects(image_np):
     results = detection_model(image_np)
 
-    # different object detection models have additional results
-    result = {key:value.numpy() for key,value in results.items()}
+    # Different object detection models may have additional results
+    result = {key: value.numpy() for key, value in results.items()}
     
     label_id_offset = 0
     image_np_with_detections = image_np.copy()
@@ -75,16 +194,14 @@ def predict2(image_np):
         category_index,
         use_normalized_coordinates=True,
         max_boxes_to_draw=200,
-        min_score_thresh=.60,
+        min_score_thresh=0.60,
         agnostic_mode=False,
-        line_thickness=2)
+        line_thickness=2
+    )
 
     result_pil_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
-    
     return result_pil_img
 
-import cv2
-
 def predict_video(video_path):
     cap = cv2.VideoCapture(video_path)
     frame_width = int(cap.get(3))
@@ -102,7 +219,7 @@ def predict_video(video_path):
         pil_image = Image.fromarray(frame)
         
         # Perform object detection on the frame
-        result_pil_img = predict(pil_image)
+        result_pil_img = predict_objects(frame)
 
         # Convert the result back to a NumPy array
         result_np_img = tf.keras.utils.img_to_array(result_pil_img)
@@ -116,24 +233,21 @@ def predict_video(video_path):
 
     return "output.avi"
 
-
-
 REPO_ID = "Louisw3399/burgerorfriesdetector"
 detection_model = load_model()
-# pil_image = Image.open(image_path)
-# image_arr = pil_image_as_numpy_array(pil_image)
-
-# predicted_img = predict(image_arr)
-# predicted_img.save('predicted.jpg')
 
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil")
-             ).launch(share=True)
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    label="Image Object Detection"
+).launch(share=True)
 
 gr.Interface(
     fn=predict_video,
     inputs=gr.Video(type="file", label="Upload a video"),
-    outputs=gr.Video(type="file", label="Download the processed video")
+    outputs=gr.Video(type="file", label="Download the processed video"),
+    label="Video Object Detection"
 ).launch(share=True)
 
+