vislearnlabpy 0.0.2.2__tar.gz → 0.0.2.3__tar.gz

{vislearnlabpy-0.0.2.2/src/vislearnlabpy.egg-info → vislearnlabpy-0.0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vislearnlabpy
-Version: 0.0.2.2
+Version: 0.0.2.3
 Summary: Visual Learning Lab utility files and pipelines
 Author-email: Tarun Sepuri <tarunsepuri@gmail.com>
 License-Expression: MIT

{vislearnlabpy-0.0.2.2 → vislearnlabpy-0.0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vislearnlabpy"
-version = "0.0.2.2"
+version = "0.0.2.3"
 authors = [
   { name="Tarun Sepuri", email="tarunsepuri@gmail.com" },
 ]

{vislearnlabpy-0.0.2.2 → vislearnlabpy-0.0.2.3}/src/vislearnlabpy/embeddings/stimuli_loader.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 import os
 import pandas as pd
 import re
@@ -12,10 +13,20 @@ import numpy as np
 
 random.seed(2)
 
+
+@dataclass
+class ImgExtractionSettings:
+    resize_dim: int = 256
+    crop_dim: int = 224
+    apply_content_crop: bool = True
+    apply_center_crop: bool = False
+    use_thumbnail: bool = False
+    change_stroke_color: bool = False
+    stroke_color: tuple = (0, 0, 0)
+    stroke_threshold: int = 200
+
+
 class ImageExtractor:
-    def __init__(self):
-        pass
-    
     @staticmethod
     def RGBA2RGB(img, background_color=(255, 255, 255)):
         """Alpha composite an RGBA Image with a specified color.
@@ -42,6 +53,41 @@ class ImageExtractor:
             img = img.convert('RGB')
         return img
 
+    @staticmethod
+    def change_stroke_color(img, target_color=(0, 0, 0), threshold=200):
+        """
+        Change the color of strokes/dark pixels in an image.
+        
+        Args:
+            img -- PIL Image object
+            target_color -- Tuple (r, g, b) for the new stroke color
+            threshold -- Brightness threshold to identify strokes (0-255)
+                        Pixels darker than this are considered strokes
+        
+        Returns:
+            PIL Image with recolored strokes
+        """
+        # Convert to RGBA if not already
+        if img.mode != 'RGBA':
+            img = img.convert('RGBA')
+        
+        # Convert to numpy array
+        img_array = np.array(img)
+        
+        # Calculate brightness (average of RGB channels)
+        brightness = img_array[:, :, :3].mean(axis=2)
+        
+        # Create mask for stroke pixels (darker than threshold)
+        stroke_mask = brightness < threshold
+        
+        # Apply new color to stroke pixels while preserving alpha
+        result = img_array.copy()
+        result[stroke_mask, 0] = target_color[0]  # R
+        result[stroke_mask, 1] = target_color[1]  # G
+        result[stroke_mask, 2] = target_color[2]  # B
+        
+        return Image.fromarray(result)
+
     @staticmethod
     def crop_to_content(img, apply_content_crop):
         """Crop image to remove white space around content."""
@@ -67,41 +113,45 @@ class ImageExtractor:
         return img
     
     @staticmethod
-    def get_transformations(resize_dim=256, crop_dim=224, apply_content_crop=True, apply_center_crop=False, use_thumbnail=False):
+    def get_transformations(settings: ImgExtractionSettings=ImgExtractionSettings()):
         """Load image transformations for dataloader.
         
         Args:
-            resize_dim: Dimension for resizing (default 256)
-            crop_dim: Dimension for center crop (default 224)
-            apply_content_crop: Whether to crop to content (default True)
-            apply_center_crop: Whether to apply center crop (default True)
-            use_thumbnail: Whether to use thumbnail resizing instead of regular resize (default False)
+            settings: ImgExtractionSettings dataclass instance
         """
         
         def combined_transform(image):
             # Step 1: Apply thumbnail resizing if requested
-            if use_thumbnail:
-                image.thumbnail((resize_dim, resize_dim), Image.Resampling.LANCZOS)
+            if settings.use_thumbnail:
+                image.thumbnail((settings.resize_dim, settings.resize_dim), Image.Resampling.LANCZOS)
 
-            # Step 2: Convert RGBA to RGB
+            # Step 2: Change stroke color if enabled (before converting to RGB)
+            if settings.change_stroke_color:
+                image = ImageExtractor.change_stroke_color(
+                    image, 
+                    settings.stroke_color, 
+                    settings.stroke_threshold
+                )
+
+            # Step 3: Convert RGBA to RGB
             img_rgb = ImageExtractor.RGBA2RGB(image)
             
-            # Step 3: Crop to content (remove whitespace) if enabled
-            img_cropped = ImageExtractor.crop_to_content(img_rgb, apply_content_crop)
+            # Step 4: Crop to content (remove whitespace) if enabled
+            img_cropped = ImageExtractor.crop_to_content(img_rgb, settings.apply_content_crop)
             return img_cropped
         
         # Build transformation pipeline
         transform_list = []
         
         # Only add regular resize if not using thumbnail
-        if not use_thumbnail:
+        if not settings.use_thumbnail:
             # resize first
-            transform_list.append(transforms.Resize(resize_dim))
+            transform_list.append(transforms.Resize(settings.resize_dim))
 
         transform_list.append(transforms.Lambda(combined_transform))
         
-        if apply_center_crop:
-            transform_list.append(transforms.CenterCrop(crop_dim))
+        if settings.apply_center_crop:
+            transform_list.append(transforms.CenterCrop(settings.crop_dim))
         
         # Add any additional transforms you might need
         #transform_list.extend([
@@ -111,6 +161,27 @@ class ImageExtractor:
         
         return transforms.Compose(transform_list)
 
+    @staticmethod
+    def save_transformed(original_file, new_file, settings: ImgExtractionSettings = None, transform=None):
+        """
+        Load, transform, and save an image.
+        
+        Args:
+            original_file: Path to input image
+            new_file: Path to save output image
+            settings: ImgExtractionSettings dataclass instance (if None, uses default)
+            transform: Custom transform (if provided, overrides settings)
+        """
+        if settings is None:
+            settings = ImgExtractionSettings()
+            
+        if transform is None:
+            transform = ImageExtractor.get_transformations(settings)
+        
+        image_data = Image.open(original_file)
+        img = transform(image_data)
+        img.save(new_file)
+
 class StimuliDataset(Dataset):
     def __init__(self, manifest, images_folder=None, id_column=None, transform=None):
         self.manifest = manifest

{vislearnlabpy-0.0.2.2 → vislearnlabpy-0.0.2.3/src/vislearnlabpy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vislearnlabpy
-Version: 0.0.2.2
+Version: 0.0.2.3
 Summary: Visual Learning Lab utility files and pipelines
 Author-email: Tarun Sepuri <tarunsepuri@gmail.com>
 License-Expression: MIT