File size: 6,226 Bytes
87520e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | import os
import shutil
import cv2
import tempfile
from flask import Flask, render_template, request, jsonify
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
import torch
# We keep gdown import for fallback
import gdown
app = Flask(__name__)
# Google Drive ID for the model (Folder)
MODEL_DRIVE_ID = "1ovFKKSc9xc1-DL9AEAi4ZSGDhiBIFt0j"
model_path = "./blip-model"
def load_model():
# Attempt to load model
try:
print(f"Loading model from {model_path}...")
processor = BlipProcessor.from_pretrained(model_path)
model = BlipForConditionalGeneration.from_pretrained(model_path)
print("Model loaded successfully.")
return processor, model
except Exception as e:
print(f"Failed to load model from {model_path}: {e}")
print("Model might not be present. Attempting download (Fallback)...")
# Fallback download logic here if needed, or we rely on Dockerfile
try:
# Check if model exists and is not empty (simple check)
if not os.path.exists(model_path) or not os.listdir(model_path):
print(f"Downloading from Google Drive...")
os.makedirs(model_path, exist_ok=True)
url = f'https://drive.google.com/drive/folders/{MODEL_DRIVE_ID}'
gdown.download_folder(url, output=model_path, quiet=False, use_cookies=False)
# Retry load
processor = BlipProcessor.from_pretrained(model_path)
model = BlipForConditionalGeneration.from_pretrained(model_path)
print("Model loaded after fallback download.")
return processor, model
except Exception as download_error:
print(f"Fallback download failed: {download_error}")
return None, None
processor, model = load_model()
# Check for GPU
if model:
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print(f"Using device: {device}")
else:
print("CRITICAL WARNING: Model not loaded. Application will start but will fail to generate captions.")
device = "cpu"
@app.route('/')
def index():
return render_template('index.html')
def generate_caption(raw_image, style):
inputs = processor(raw_image, return_tensors="pt").to(device)
gen_kwargs = {
"max_new_tokens": 100,
"min_length": 20,
"num_beams": 5
}
if style == 'detailed':
gen_kwargs.update({
"min_length": 40,
"max_new_tokens": 150,
"repetition_penalty": 1.2,
"num_beams": 5
})
elif style == 'concise':
gen_kwargs.update({
"max_new_tokens": 20,
"min_length": 5,
"num_beams": 3
})
elif style == 'creative':
gen_kwargs.update({
"do_sample": True,
"top_k": 50,
"top_p": 0.92,
"temperature": 0.8,
"num_beams": 3,
"repetition_penalty": 1.3,
"min_length": 25
})
elif style == 'hashtags':
gen_kwargs.update({
"max_new_tokens": 30,
"min_length": 5,
"num_beams": 3
})
out = model.generate(**inputs, **gen_kwargs)
caption = processor.decode(out[0], skip_special_tokens=True)
if style == 'hashtags':
words = caption.split()
hashtags = []
for w in words:
clean_w = ''.join(e for e in w if e.isalnum())
if len(clean_w) > 2:
hashtags.append(f"#{clean_w.lower()}")
if not hashtags:
hashtags = ["#image", "#photo"]
caption = " ".join(hashtags)
return caption
def process_video(file, style):
temp_dir = tempfile.mkdtemp()
temp_video_path = os.path.join(temp_dir, file.filename)
file.save(temp_video_path)
captions = []
try:
cap = cv2.VideoCapture(temp_video_path)
fps = round(cap.get(cv2.CAP_PROP_FPS))
if fps == 0:
fps = 30 # fallback
frame_count = 0
success = True
while success:
success, frame = cap.read()
if not success:
break
# Extract 1 frame per second
if frame_count % fps == 0:
timestamp_sec = frame_count // fps
# Convert BGR to RGB
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(frame_rgb)
caption = generate_caption(pil_img, style)
captions.append({
"timestamp": timestamp_sec,
"caption": caption
})
frame_count += 1
cap.release()
finally:
# Cleanup
if os.path.exists(temp_video_path):
os.remove(temp_video_path)
os.rmdir(temp_dir)
return jsonify({'captions': captions, 'type': 'video'})
@app.route('/caption', methods=['POST'])
def caption_image():
if not model:
return jsonify({'error': 'Model not loaded'}), 500
file = request.files.get('file') or request.files.get('image')
if not file or file.filename == '':
return jsonify({'error': 'No selected file or image provided'}), 400
try:
style = request.form.get('style', 'default')
mimetype = file.mimetype
if mimetype and mimetype.startswith('video/'):
return process_video(file, style)
else:
raw_image = Image.open(file.stream).convert('RGB')
caption = generate_caption(raw_image, style)
return jsonify({'caption': caption, 'type': 'image'})
except Exception as e:
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
# Run on 0.0.0.0 and port 7860 for Hugging Face Spaces
app.run(host='0.0.0.0', port=7860)
|