Update app.py

app.py

@@ -1,113 +1,175 @@
-import cv2
+import gradio as gr
 import torch
-from PIL import Image
+import cv2
 import numpy as np
-import os
-import shutil
-import gradio as gr
-from transformers import LlavaNextVideoProcessor, LlavaNextVideoForConditionalGeneration, BitsAndBytesConfig
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_id = "llava-hf/LLaVA-NeXT-Video-7B-hf"
+from PIL import Image
+from transformers import AutoProcessor, LlavaNextForConditionalGeneration
 
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4"
-)
+MODEL_ID = "arjunanand13/gas_pipe_llava_finetunedv2"
 
-model = LlavaNextVideoForConditionalGeneration.from_pretrained(
-    model_id,
-    quantization_config=quantization_config,
-    low_cpu_mem_usage=True,
-    device_map="auto"
-)
+@torch.no_grad()
+def load_model():
+    processor = AutoProcessor.from_pretrained(MODEL_ID)
+    model = LlavaNextForConditionalGeneration.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    return processor, model
 
-processor = LlavaNextVideoProcessor.from_pretrained(model_id)
+processor, model = load_model()
 
-def sample_frames(video_path, num_frames):
-    output_dir = "/tmp/processed_frames"
+def extract_frames_from_video(video_path, num_frames=4):
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise ValueError(f"Cannot open video: {video_path}")
     
-    if os.path.exists(output_dir):
-        shutil.rmtree(output_dir)
-    os.makedirs(output_dir)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    frame_indices = np.linspace(0, max(0, total_frames - 1), num_frames, dtype=int)
     
-    video = cv2.VideoCapture(video_path)
-    
-    if not video.isOpened():
-        raise ValueError(f"Could not open video file: {video_path}")
-    
-    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
-    interval = max(1, total_frames // num_frames)
     frames = []
-    frame_count = 0
+    for frame_idx in frame_indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+        ret, frame = cap.read()
+        if ret:
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame_pil = Image.fromarray(frame_rgb)
+            frame_resized = frame_pil.resize((112, 112), Image.Resampling.LANCZOS)
+            frames.append(frame_resized)
     
-    for i in range(total_frames):
-        ret, frame = video.read()
-        if not ret:
-            continue
-        if i % interval == 0 and len(frames) < num_frames:
-            cv2.imwrite(f"{output_dir}/frame_{frame_count:03d}.jpg", frame)
-            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-            frames.append(pil_img)
-            frame_count += 1
+    cap.release()
     
-    video.release()
+    while len(frames) < 4:
+        if frames:
+            frames.append(frames[-1].copy())
+        else:
+            frames.append(Image.new('RGB', (112, 112), color='black'))
     
-    frame_paths = [f"{output_dir}/frame_{i:03d}.jpg" for i in range(frame_count)]
+    grid_image = Image.new('RGB', (224, 224))
+    positions = [(0, 0), (112, 0), (0, 112), (112, 112)]
     
-    return frames, frame_paths
+    for i, frame in enumerate(frames[:4]):
+        grid_image.paste(frame, positions[i])
+    
+    return grid_image
 
-def analyze_video(video_path):
-    conversation = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": "Analyze this gas pipe quality control video. Answer these two questions with True/False: 1) DIPPED IN WATER: Was the pipe dipped in water for testing? Look for pipe being submerged in water container. 2) BUBBLES AFTER IMMERSION: After the pipe was fully immersed (ignore initial displacement bubbles), were there any bubbles indicating a leak? Format: DIPPED IN WATER: True/False, BUBBLES AFTER IMMERSION: True/False"},
-                {"type": "video"},
-            ],
-        },
-    ]
+def predict_cheating(video_file):
+    if video_file is None:
+        return "Please upload a video file", None
     
-    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+    try:
+        grid_image = extract_frames_from_video(video_file, num_frames=4)
+        
+        prompt = """nalyze the gas‐pipe QC video and reply ONLY with valid JSON:
+        {"cheating": true} if any step is skipped, else {"cheating": false}.
+        
+        Required steps (no cheating):
+        1. Connect pipe to gas supply  
+        2. Submerge in water and check for bubbles  
+        3. Mark pipe and move to designated area  
+        
+        Cheating (any violation):
+        - Skipping water submersion or bubble check  
+        - Moving pipe without marking/testing  
+        
+        Examples:
+        - Dipped, checked bubbles, marked → {"cheating": false}  
+        - Moved without testing → {"cheating": true}  """
+        inputs = processor(text=prompt, images=grid_image, return_tensors="pt")
+        
+        generated_ids = model.generate(
+            **inputs,
+            max_new_tokens=50,
+            do_sample=False,
+            temperature=0.1
+        )
+        
+        result = processor.decode(generated_ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
+        
+        if "true" in result.lower():
+            prediction = "CHEATING DETECTED"
+            explanation = "The system detected violations in the testing procedure."
+        else:
+            prediction = "COMPLIANT PROCEDURE"
+            explanation = "The testing procedure appears to follow proper protocols."
+        
+        return f"{prediction}\n\n{explanation}\n\nRaw output: {result}", grid_image
+        
+    except Exception as e:
+        return f"Error processing video: {str(e)}", None
+
+with gr.Blocks(title="Gas Pipe Quality Control Detection") as demo:
+    gr.Markdown("# Gas Pipe Quality Control Detection")
+    gr.Markdown("Upload a video of gas pipe testing to detect compliance violations.")
     
-    video_frames, frame_paths = sample_frames(video_path, 20)
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            ### Proper Procedure (No Cheating)
+            1. Connect pipe to gas supply
+            2. Immerse pipe in water container
+            3. Check for bubbles (bubbles = leak)
+            4. Mark pipe with marker/pen
+            5. Move tested pipe to designated area
+            """)
+            
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            ### Cheating Behaviors
+            - Skipping water immersion test
+            - Moving pipe directly without testing
+            - Not checking for bubbles properly
+            - Bypassing any required step
+            """)
     
-    inputs = processor(text=prompt, videos=video_frames, padding=True)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    with gr.Row():
+        with gr.Column(scale=2):
+            video_input = gr.Video(
+                label="Upload Gas Pipe Testing Video",
+                height=300
+            )
+            
+            analyze_btn = gr.Button(
+                "Analyze Video", 
+                variant="primary",
+                size="lg"
+            )
+        
+        with gr.Column(scale=2):
+            result_text = gr.Textbox(
+                label="Detection Result",
+                lines=8,
+                max_lines=10
+            )
+            
+            processed_image = gr.Image(
+                label="Processed Video Frames (2x2 Grid)",
+                height=300
+            )
     
-    output = model.generate(
-        **inputs,
-        max_new_tokens=100,
-        do_sample=True,
-        temperature=0.3,
-        top_p=0.9,
-        top_k=50,
-        repetition_penalty=1.1,
-        pad_token_id=processor.tokenizer.eos_token_id
+    analyze_btn.click(
+        fn=predict_cheating,
+        inputs=[video_input],
+        outputs=[result_text, processed_image]
     )
     
-    result = processor.decode(output[0][2:], skip_special_tokens=True)
+    gr.Markdown("### Example Cases")
     
-    return frame_paths, result
-
-examples = [
-    ["07.mp4"],
-    ["09.mp4"], 
-    ["29.mp4"]
-]
-
-iface = gr.Interface(
-    fn=analyze_video,
-    inputs=gr.Video(),
-    outputs=[
-        gr.Gallery(label="Processed Frames"),
-        gr.Textbox(label="LLM Analysis", lines=10)
-    ],
-    title="Gas Pipe Quality Control Analyzer",
-    examples=examples,
-    cache_examples=False
-)
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            **Cheating Examples:**
+            """)
+            
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            **Compliant Examples:**
+            """)
+    
+    gr.Markdown("""
+    ---
+    Model: Fine-tuned LLaVA v1.6 Mistral 7B | Repository: arjunanand13/gas_pipe_llava_finetunedv2
+    """)
 
-iface.launch(share=True)
+if __name__ == "__main__":
+    demo.launch(share=True)