app.py

import gradio as gr
import io
import numpy as np
import torch
from decord import cpu, VideoReader, bridge
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import BitsAndBytesConfig

MODEL_PATH = "THUDM/cogvlm2-llama3-caption"
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16


def get_step_info(step_name):
    """Returns detailed information about a manufacturing step."""
    step_details = {
        "Step 1": {
            "Name": "Bead Insertion",
            "Standard Time": "4 seconds",
            "Analysis": "Observe the bead placement process. If the insertion exceeds 4 seconds, identify potential issues such as missing beads, technician errors, or machinery malfunction."
        },
        "Step 2": {
            "Name": "Inner Liner Apply",
            "Standard Time": "4 seconds",
            "Analysis": "Check for manual intervention during the inner layer application. If adjustment is required, it may indicate improper alignment or issues with the layer material."
        },
        "Step 3": {
            "Name": "Ply1 Apply",
            "Standard Time": "4 seconds",
            "Analysis": "Determine if the technician is manually adjusting the first ply. Manual intervention might suggest improper ply placement or machine misalignment."
        },
        "Step 4": {
            "Name": "Bead Set",
            "Standard Time": "8 seconds",
            "Analysis": "Observe the bead setting process. Delays may result from bead misalignment, machine pauses, or lack of technician involvement."
        },
        "Step 5": {
            "Name": "Turnup",
            "Standard Time": "4 seconds",
            "Analysis": "Examine the turnup step for any technician involvement or pauses in machine operation. Reasons for delays might include material misalignment or equipment issues."
        },
        "Step 6": {
            "Name": "Sidewall Apply",
            "Standard Time": "14 seconds",
            "Analysis": "If a technician is repairing the sidewall, this may indicate material damage or improper initial application. Look for signs of excessive manual handling."
        },
        "Step 7": {
            "Name": "Sidewall Stitching",
            "Standard Time": "5 seconds",
            "Analysis": "Observe the stitching process. Delays could occur due to machine speed inconsistencies or technician intervention for correction."
        },
        "Step 8": {
            "Name": "Carcass Unload",
            "Standard Time": "7 seconds",
            "Analysis": "Ensure a technician is present for the carcass unload. If absent, note their return time and identify potential reasons for their absence."
        }
    }
    
    return step_details.get(step_name, {"Error": "Invalid step name. Please provide a valid step number."})


def load_video(video_data, strategy='chat'):
    """Loads and processes video data into a format suitable for model input."""
    bridge.set_bridge('torch')
    num_frames = 24
    
    if isinstance(video_data, str): 
        decord_vr = VideoReader(video_data, ctx=cpu(0))
    else:  
        decord_vr = VideoReader(io.BytesIO(video_data), ctx=cpu(0))
    
    frame_id_list = []
    total_frames = len(decord_vr)
    timestamps = [i[0] for i in decord_vr.get_frame_timestamp(np.arange(total_frames))]
    max_second = round(max(timestamps)) + 1
    
    for second in range(max_second):
        closest_num = min(timestamps, key=lambda x: abs(x - second))
        index = timestamps.index(closest_num)
        frame_id_list.append(index)
        if len(frame_id_list) >= num_frames:
            break

    video_data = decord_vr.get_batch(frame_id_list)
    video_data = video_data.permute(3, 0, 1, 2)
    return video_data

def load_model():
    """Loads the pre-trained model and tokenizer with quantization configurations."""
    quantization_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=TORCH_TYPE,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4"
    )
    
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        torch_dtype=TORCH_TYPE,
        trust_remote_code=True,
        quantization_config=quantization_config,
        device_map="auto"
    ).eval()
    
    return model, tokenizer

def predict(prompt, video_data, temperature, model, tokenizer):
    """Generates predictions based on the video and textual prompt."""
    video = load_video(video_data, strategy='chat')
    
    inputs = model.build_conversation_input_ids(
        tokenizer=tokenizer,
        query=prompt,
        images=[video],
        history=[],
        template_version='chat'
    )
    
    inputs = {
        'input_ids': inputs['input_ids'].unsqueeze(0).to(DEVICE),
        'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to(DEVICE),
        'attention_mask': inputs['attention_mask'].unsqueeze(0).to(DEVICE),
        'images': [[inputs['images'][0].to(DEVICE).to(TORCH_TYPE)]],
    }
    
    gen_kwargs = {
        "max_new_tokens": 2048,
        "pad_token_id": 128002,
        "top_k": 1,
        "do_sample": False,
        "top_p": 0.1,
        "temperature": temperature,
    }
    
    with torch.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return response

def get_analysis_prompt(step_number):
    """Constructs the prompt for analyzing manufacturing delays based on the selected step."""
    return f"""You are an AI expert system specializing in manufacturing processes. 
Your task is to analyze video footage from Step {step_number} of a tire manufacturing process and identify any issues based on the observed footage.
- Focus on identifying signs of delay or disruption.
- If no person is visible, it may indicate a staffing issue.
- If a person is seen modifying the tire, they may be repairing defects or handling material issues.
- Carefully examine for mechanical failures, material problems, or human involvement.

Provide an analysis of the video by determining the most likely cause of delay in this step, and explain why this conclusion was reached based on the visual evidence."""

model, tokenizer = load_model()

def inference(video, step_number):
    """Analyzes video to predict possible issues based on the manufacturing step."""
    try:
        if not video:
            return "Please upload a video first."
        
        prompt = get_analysis_prompt(step_number)
        temperature = 0.8
        response = predict(prompt, video, temperature, model, tokenizer)
        
        return response
    except Exception as e:
        return f"An error occurred during analysis: {str(e)}"

def create_interface():
    """Creates the Gradio interface for the Manufacturing Analysis System."""
    with gr.Blocks() as demo:
        gr.Markdown("""
        # Manufacturing Analysis System
        Upload a video of the manufacturing step and select the step number. 
        The system will analyze the video and provide observations.
        """)
        
        with gr.Row():
            with gr.Column():
                video = gr.Video(label="Upload Manufacturing Video", sources=["upload"])
                step_number = gr.Dropdown(
                    choices=[f"Step {i}" for i in range(1, 9)],
                    label="Manufacturing Step"
                )
                analyze_btn = gr.Button("Analyze", variant="primary")
            
            with gr.Column():
                output = gr.Textbox(label="Analysis Result", lines=10)
        
        gr.Examples(
            examples=[
                ["7838_step2_2_eval.mp4", "Step 2"],
                ["7838_step6_2_eval.mp4", "Step 6"]
            ],
            inputs=[video, step_number],
            cache_examples=False
        )
        
        analyze_btn.click(
            fn=inference,
            inputs=[video, step_number],
            outputs=[output]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.queue().launch(share=True)