Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import os
 import whisper
 import cv2
+import io
+from PIL import Image
 import json
 import tempfile
 import torch
@@ -87,13 +89,13 @@ class VideoClassifier:
         start = time.time()
         mp4_file = video_input
         video_name = mp4_file.split("/")[-1]
-        wav_file = "audiotrack.wav"
+        wav_file = "results/audiotrack.wav"
         video_clip = VideoFileClip(mp4_file)
         audioclip = video_clip.audio
         wav_file = audioclip.write_audiofile(wav_file)
         audioclip.close()
         video_clip.close()
-        audiotrack = "audiotrack.wav"
+        audiotrack = "results/audiotrack.wav"
         result = self.whisper_model.transcribe(audiotrack, fp16=False)
         transcript = result["text"]
         print("TRANSCRIPT",transcript)
@@ -102,35 +104,33 @@ class VideoClassifier:
         time_taken_1 = round(end - start, 3)
         # print("Time taken from video to transcript:", time_taken_1)
 
-        
         video = cv2.VideoCapture(video_input)
         length = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-        # print(f'There are {length} frames in this video')
-        # args = self.parser.parse_args()
         no_of_frame = int(self.no_of_frames)
-        temp_div = length//(no_of_frame)
+        temp_div = length // no_of_frame
         currentframe = 50
         caption_text = []
+
         for i in range(no_of_frame):
-            # reading from frame
             video.set(cv2.CAP_PROP_POS_FRAMES, currentframe)
-            ret,frame = video.read()
+            ret, frame = video.read()
             if ret:
-                content = self.img_cap.predict_step([frame])
-                print("content",content)
+        
+                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                image = Image.fromarray(frame)
+                img_byte_arr = io.BytesIO()
+                image.save(img_byte_arr, format='JPEG')  # Save as JPEG or any other format your model supports
+                img_byte_arr.seek(0)
+                
+                content = self.img_cap.predict_from_memory([img_byte_arr])
+                print("content", content)
                 caption_text.append(content[0])
-                currentframe += temp_div-1
+                currentframe += temp_div - 1
             else:
                 break
-        def listToString(s):
-            str1 = ", "
-            return (str1.join(s))
-        captions = listToString(caption_text)
-        print("CAPTIONS",captions)
 
-        end1 = time.time()
-        time_taken_2 = round(end1 - end, 3)
-        # print("Time taken from transcript to image_caption:", time_taken_2)
+        captions = ", ".join(caption_text)
+        print("CAPTIONS", captions)
         video.release()
         cv2.destroyAllWindows()
 
@@ -142,23 +142,18 @@ class VideoClassifier:
 
 
         template1 = '''Given below are the different type of main video classes  
-
         {main_categories}
-
         You are a text classifier that catergorises the transcript and captions into one main class whose context match with one main class and only generate main class name no need of sub classe or explanation.
         Give more importance to Transcript while classifying .
         Transcript: {transcript}
-
         Captions: {captions}
-
         Return only the answer chosen from list and nothing else
         Main-class =>  '''
 
         prompt1 = PromptTemplate(template=template1, input_variables=['main_categories', 'transcript', 'captions'])
         print("PROMPT 1",prompt1)
         prompt_text = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
-        start2 = time.time()
-        time_taken_3 = round(start2 - end1, 3)
+      
         response = self.genai_model.generate_content(contents=prompt_text)
         main_class = response.text
         
@@ -203,15 +198,11 @@ class VideoClassifier:
         sub_categories = Path("sub_categories.txt").read_text()
 
         template2 = '''Given below are the sub classes of {main_class}.
-
         {sub_categories}
-
         You are a text classifier that catergorises the transcript and captions into one sub class whose context match with one sub class and only generate sub class name, Don't give explanation .
         Give more importance to Transcript while classifying .
         Transcript: {transcript}
-
         Captions: {captions}
-
         Return only the Sub-class answer chosen from list and nothing else
         Answer in the format:
         Main-class => {main_class}
@@ -220,14 +211,10 @@ class VideoClassifier:
         
         prompt2 = PromptTemplate(template=template2, input_variables=['sub_categories', 'transcript', 'captions','main_class'])
         prompt_text2 = template1.format(main_categories=main_categories, transcript=transcript, captions=captions)
-        start2 = time.time()
-        time_taken_3 = round(start2 - end1, 3)
         response = self.genai_model.generate_content(contents=prompt_text2)
         sub_class = response.text
         print("Preprocess Answer",sub_class)
-        end2 = time.time()
-        time_taken_predict = round(end2 - start2, 3)
-        time_taken_total = round(end2 - start, 3)
+        
         # print("Time taken by model to predict:", time_taken_predict)
         # print("Total time taken:", time_taken_total)