Update app.py

app.py

@@ -1,79 +1,107 @@
-MODEL_NAME="allenai/MolmoE-1B-0924"
-from transformers import AutoModelForCausalLM
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)    
 from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 from PIL import Image
-import requests
+import torch
+import spaces
 
-# load the processor
+# Load the processor and model
 processor = AutoProcessor.from_pretrained(
-    'allenai/MolmoE-1B-0924',
+    'allenai/Molmo-1B-D-0924',
     trust_remote_code=True,
     torch_dtype='auto',
     device_map='auto'
 )
 
-# load the model
 model = AutoModelForCausalLM.from_pretrained(
-    'allenai/MolmoE-1B-0924',
+    'allenai/Molmo-1B-D-0924',
     trust_remote_code=True,
     torch_dtype='auto',
     device_map='auto'
 )
 
-# process the image and text
-inputs = processor.process(
-    images=[Image.open(requests.get("https://picsum.photos/id/237/536/354", stream=True).raw)],
-    text="Describe this image."
-)
 
-# move inputs to the correct device and make a batch of size 1
-inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+@spaces.GPU(duration=120)
+def process_image_and_text(image, text):
+    # Process the image and text
+    inputs = processor.process(
+        images=[Image.fromarray(image)],
+        text=text
+    )
 
-# generate output; maximum 200 new tokens; stop generation when <|endoftext|> is generated
-output = model.generate_from_batch(
-    inputs,
-    GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
-    tokenizer=processor.tokenizer
-)
+    # Move inputs to the correct device and make a batch of size 1
+    inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}
+
+    # Generate output
+    output = model.generate_from_batch(
+        inputs,
+        GenerationConfig(max_new_tokens=200, stop_strings="<|endoftext|>"),
+        tokenizer=processor.tokenizer
+    )
 
-# only get generated tokens; decode them to text
-generated_tokens = output[0,inputs['input_ids'].size(1):]
-generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+    # Only get generated tokens; decode them to text
+    generated_tokens = output[0, inputs['input_ids'].size(1):]
+    generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
-# print the generated text
-print(generated_text)
+    return generated_text
 
 # >>> This photograph captures a small black puppy, likely a Labrador or a similar breed,
 #     sitting attentively on a weathered wooden deck. The deck, composed of three...
 
 
-# import cv2
+import cv2
 
 
-# class Solution():
-#     def __init__(self,prompt):
-#         self.prompt= prompt
-#         self.output_dir=None
+class Video():
+    def __init__(self,prompt):
+        self.prompt= prompt
+        self.output_dir=None
         
-#     # read a mp4 file and getting its frame at a particular interval.
-#     def read_frame(self,file,interval=1):
-#         video=cv2.VideoCapture(file)
-#         fps= video.get(cv2.CAP_PROP_FPS)
-
-#         frame_interval= fps*interval# fps= 24 frame/sec and interval = 1 sec so frame interval = 24 frame         
-#         while True:
-#             success, frame=video.read()
-#             if not success:
-#                 break
-
-#             if frame % frame_interval==0:
-#                 # process this frame
-#                 """
-#                 .. to do 
-#                 """
+    # read a mp4 file and getting its frame at a particular interval.
+    def read_frame(self,file,interval=1):
+        video=cv2.VideoCapture(file)
+        fps= video.get(cv2.CAP_PROP_FPS)
+
+        frame_interval= fps*interval# fps= 24 frame/sec and interval = 1 sec so frame interval = 24 frame         
+        while True:
+            success, frame=video.read()
+            if not success:
+                break
+
+            if frame % frame_interval==0:
+                # process this frame
+                """
+                send the frame to MolMO which will return either co-ordinates : x,y or Null
+                """
+                self.find(frame)
                 
-#     def find(self,input_message):
+    def find(self,frame):
+        """
+        This function will take in the frame and input message and point to all the messages in the picture.
+        """
+        model= Model()
+        text = model.generate(frame,self.prompt)
+        x,y=extract_coordinates(text)
+        if (x,y):
+            annotate_the_image_with_pointer(x,y,frame)
+        else:
+            """
+            read next frame
+            """
+            break;
+
+    import re
+    def extract_coordinates(text):
+        object=re.compile(r"\(([^)]+)\)")
+        co_ord= object.search(text)
+        if co_ord:
+             # Split the captured text on the comma to get the x and y values
+            x, y = map(float, co_ord.group(1).split(','))
+        coordinates = (x, y)
+        
+        return co_ord        
+
+    def annotate_the_image_with_pointer(x,y,frame):
+        return cv2.circle(frame,(x,y),2,(255,0,0),2)
+        
         
         
                 
@@ -81,8 +109,8 @@ print(generated_text)
 
 
 
-    # read a .mp4 file
-    # get a interval N spaced
+    read a .mp4 file
+    get a interval N spaced