Apollo_GenAI / app.py
VishalD1234's picture
Update app.py
73a629c verified
raw
history blame
10.1 kB
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-video-llama3-chat"
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_number):
"""Returns detailed information about a manufacturing step."""
step_details = {
1: {
"Name": "Bead Insertion",
"Standard Time": "4 seconds",
"Video_substeps_expected": {
"0-1 second": "Machine starts bead insertion process.",
"1-3 seconds": "Beads are aligned and positioned.",
"3-4 seconds": "Final adjustment and confirmation of bead placement."
}
},
2: {
"Name": "Inner Liner Apply",
"Standard Time": "4 seconds",
"Video_substeps_expected": {
"0-1 second": "Machine applies the first layer of the liner.",
"1-3 seconds": "Technician checks alignment and adjusts if needed.",
"3-4 seconds": "Final inspection and confirmation."
}
},
3: {
"Name": "Ply1 Apply",
"Standard Time": "4 seconds",
"Video_substeps_expected": {
"0-2 seconds": "First ply is loaded onto the machine.",
"2-4 seconds": "Technician inspects and adjusts ply placement."
}
},
4: {
"Name": "Bead Set",
"Standard Time": "8 seconds",
"Video_substeps_expected": {
"0-3 seconds": "Bead is positioned and pre-set.",
"3-6 seconds": "Machine secures the bead in place.",
"6-8 seconds": "Technician confirms the bead alignment."
}
},
5: {
"Name": "Turnup",
"Standard Time": "4 seconds",
"Video_substeps_expected": {
"0-2 seconds": "Turnup process begins with machine handling.",
"2-4 seconds": "Technician inspects the turnup and makes adjustments if necessary."
}
},
6: {
"Name": "Sidewall Apply",
"Standard Time": "14 seconds",
"Video_substeps_expected": {
"0-5 seconds": "Sidewall material is positioned by the machine.",
"5-10 seconds": "Technician checks for alignment and begins application.",
"10-14 seconds": "Final adjustments and confirmation of sidewall placement."
}
},
7: {
"Name": "Sidewall Stitching",
"Standard Time": "5 seconds",
"Video_substeps_expected": {
"0-2 seconds": "Stitching process begins automatically.",
"2-4 seconds": "Technician inspects stitching for any irregularities.",
"4-5 seconds": "Machine completes stitching process."
}
},
8: {
"Name": "Carcass Unload",
"Standard Time": "7 seconds",
"Video_substeps_expected": {
"0-3 seconds": "Technician unloads(removes) carcass(tire) from the machine."
}
}
}
return step_details.get(step_number, {"Error": "Invalid step number. 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, possible_reasons):
"""Constructs the prompt for analyzing delay reasons based on the selected step."""
return f"""You are an AI expert system specialized in analyzing manufacturing processes and identifying production delays in tire manufacturing. Your role is to accurately classify delay reasons based on visual evidence from production line footage.
Task Context:
You are analyzing video footage from Step {step_number} of a tire manufacturing process where a delay has been detected. Your task is to determine the most likely cause of the delay from the following possible reasons:
{', '.join(possible_reasons)}
{get_step_info(step_number)}
Required Analysis:
Carefully observe the video for visual cues indicating production interruption.
If no person is visible in any of the frames, the reason probably might be due to his absence.
If a person is visible in the video and is observed touching and modifying the layers of the tire, it means there is a issue with tyre being patched hence he is repairing it.
Compare observed evidence against each possible delay reason.
Select the most likely reason based on visual evidence.
Please provide your analysis in the following format:
1. Selected Reason: [State the most likely reason from the given options]
2. Visual Evidence: [Describe specific visual cues that support your selection]
3. Reasoning: [Explain why this reason best matches the observed evidence]
4. Alternative Analysis: [Brief explanation of why other possible reasons are less likely]
Important: Base your analysis solely on visual evidence from the video. Focus on concrete, observable details rather than assumptions. Clearly state if no person or specific activity is observed."""
model, tokenizer = load_model()
def inference(video, step_number):
"""Analyzes video to predict the most likely cause of delay in the selected manufacturing step."""
try:
if not video:
return "Please upload a video first."
possible_reasons = DELAY_REASONS[step_number]
prompt = get_analysis_prompt(step_number, possible_reasons)
temperature = 0.3
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"],
["7838_step8_1_eval.mp4", "Step 8"],
["7993_step6_3_eval.mp4", "Step 6"],
["7993_step8_3_eval.mp4", "Step 8"]
],
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)