0xalivecow/senju
1
0
Fork 0
mirror of https://github.com/senju1337/senju.git synced 2025-12-24 07:39:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

chore: fix flake8 warnings

Browse source 
Refs: OPS-85
This commit is contained in:
Christoph J. Scherr 2025-03-23 18:53:01 +01:00
parent 6dd8400b61
commit 021402bf10
No known key found for this signature in database
GPG key ID: 9EB784BB202BB7BB
1 changed files with 7 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

11
senju/image_reco.py
View file

@ -7,10 +7,12 @@ from transformers import BlipProcessor, BlipForConditionalGeneration
class ImageDescriptionGenerator: class ImageDescriptionGenerator:
def __init__(self, model_name="Salesforce/blip-image-captioning-base"): def __init__(self, model_name="Salesforce/blip-image-captioning-base"):
""" """
Initialize an image description generator using a vision-language model. Initialize an image description generator using a vision-language
model.
Args: Args:
model_name: The name of the model to use (default: BLIP captioning model) model_name: The name of the model to use
(default: BLIP captioning model)
""" """
self.device = "cuda" if torch.cuda.is_available() else "cpu" self.device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {self.device}") print(f"Using device: {self.device}")
@ -27,7 +29,8 @@ class ImageDescriptionGenerator:
max_length: Maximum length of the generated caption max_length: Maximum length of the generated caption
Returns: Returns:
dict: A dictionary containing the generated description and confidence score dict: A dictionary containing the generated description
and confidence score
""" """
# Convert uploaded bytes to image # Convert uploaded bytes to image
img = Image.open(io.BytesIO(image_data)).convert("RGB") img = Image.open(io.BytesIO(image_data)).convert("RGB")
@ -52,7 +55,7 @@ class ImageDescriptionGenerator:
return { return {
"description": caption, "description": caption,
"confidence": None # Most caption models don't provide confidence scores "confidence": None
} }