Upload FlareDetectionPipeline

config.json

@@ -1,5 +1,5 @@
 {
-  "_name_or_path": "results/fcn4flare-kepler_flare/checkpoint-2085",
+  "_name_or_path": "Maxwell-Jia/fcn4flare",
   "architectures": [
     "FCN4FlareModel"
   ],
@@ -7,6 +7,15 @@
     "AutoConfig": "Maxwell-Jia/fcn4flare--configuration_fcn4flare.FCN4FlareConfig",
     "AutoModel": "Maxwell-Jia/fcn4flare--modeling_fcn4flare.FCN4FlareModel"
   },
+  "custom_pipelines": {
+    "flare_detection": {
+      "impl": "flare_detection.FlareDetectionPipeline",
+      "pt": [
+        "AutoModel"
+      ],
+      "tf": []
+    }
+  },
   "depth": 10,
   "dilation": [
     1,

flare_detection.py

@@ -0,0 +1,233 @@
+from typing import Dict, List, Union
+import numpy as np
+import torch
+from transformers import Pipeline
+from astropy.io import fits
+
+class FlareDetectionPipeline(Pipeline):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.call_count = 0
+        
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {}
+        postprocess_kwargs = {}
+        
+        # Add parameters that need to be passed to specific steps
+        return preprocess_kwargs, {}, postprocess_kwargs
+
+    def preprocess(self, light_curve: Union[np.ndarray, str, List[str]], **kwargs) -> Dict[str, torch.Tensor]:
+        """Preprocess the input light curve from FITS files.
+        
+        Args:
+            light_curve: Single FITS file path, list of FITS file paths, or numpy array
+        """
+        # Convert single path to list
+        if isinstance(light_curve, str):
+            light_curve = [light_curve]
+        
+        # Handle list of FITS file paths
+        if isinstance(light_curve, list) and isinstance(light_curve[0], str):
+            # Read data from all FITS files
+            flux_data = []
+            times_data = []
+            lengths = []  # Store lengths of each light curve
+            
+            # First pass: get max length and collect data
+            max_length = 0
+            for fits_path in light_curve:
+                with fits.open(fits_path) as hdul:
+                    time = hdul[1].data['TIME'].astype(np.float32)
+                    flux = hdul[1].data['PDCSAP_FLUX'].astype(np.float32)
+                    # Normalize flux
+                    flux = flux / np.nanmedian(flux)
+                    
+                    max_length = max(max_length, len(flux))
+                    lengths.append(len(flux))
+                    flux_data.append(flux)
+                    times_data.append(time)
+            
+            # Second pass: pad sequences
+            padded_flux = []
+            padded_times = []
+            sequence_mask = []
+            
+            for flux, time, length in zip(flux_data, times_data, lengths):
+                # Create padding
+                pad_length = max_length - length
+                
+                # Pad flux and time arrays
+                padded_f = np.pad(flux, (0, pad_length), mode='constant', constant_values=np.nan)
+                padded_t = np.pad(time, (0, pad_length), mode='constant', constant_values=np.nan)
+                
+                # Create mask (1 for real values, 0 for padding)
+                mask = np.ones(length)
+                mask = np.pad(mask, (0, pad_length), mode='constant', constant_values=0)
+                
+                padded_flux.append(padded_f)
+                padded_times.append(padded_t)
+                sequence_mask.append(mask)
+            
+            # Store time data as attribute for use in postprocessing
+            self.time_series = np.array(padded_times)
+            # Convert to arrays
+            flux_array = np.array(padded_flux)
+            sequence_mask = np.array(sequence_mask)
+            
+            # Add channel dimension
+            flux_array = flux_array.reshape(flux_array.shape[0], flux_array.shape[1], 1)
+            
+            # Convert to torch tensors
+            inputs = torch.tensor(flux_array, dtype=torch.float32)
+            mask = torch.tensor(sequence_mask, dtype=torch.float32)
+            
+        return {
+            "input_features": inputs,
+            "sequence_mask": mask
+        }
+
+    def _forward(self, model_inputs, **forward_params):
+        """Forward pass through the model.
+        
+        Args:
+            model_inputs: Dictionary containing input tensors
+            forward_params: Additional parameters for the forward pass
+        """
+        if model_inputs is None:
+            raise ValueError("model_inputs cannot be None. Check if preprocess method is returning correct dictionary.")
+        
+        if "input_features" not in model_inputs:
+            raise KeyError("model_inputs must contain 'input_features' key.")
+        
+        # Save input_features for use in postprocessing
+        self.input_features = model_inputs["input_features"]
+        
+        # Ensure input_features is properly passed to the model
+        return self.model(
+            input_features=model_inputs["input_features"],
+            sequence_mask=model_inputs.get("sequence_mask", None),
+            return_dict=True
+        )
+
+    def postprocess(self, model_outputs, **kwargs):
+        """
+        Postprocess the model outputs to detect flare events.
+        Returns a list of dictionaries containing flare events information.
+        """
+        logits = model_outputs.logits
+        predictions = torch.sigmoid(logits).squeeze(-1)
+        binary_predictions = (predictions > 0.5).long()
+        
+        # Convert to numpy for processing
+        predictions_np = binary_predictions.cpu().numpy()
+        flux_data = self.input_features.cpu().numpy()
+        
+        flare_events = []
+        
+        def is_valid_flare(flux, start_idx, end_idx, peak_idx):
+            """Helper function to validate flare events
+            
+            Args:
+                flux: Array of flux values
+                start_idx: Start index of potential flare
+                end_idx: End index of potential flare
+                peak_idx: Peak index of potential flare
+                
+            Returns:
+                bool: True if the event is a valid flare, False otherwise
+            """
+            # Duration of a flare should be longer than 3 cadences
+            if end_idx - start_idx < 2:
+                return False
+            
+            try:
+                # If start time is the peak time, flux[start] must be greater than flux[start-1]
+                if peak_idx == start_idx and flux[peak_idx] <= flux[peak_idx - 1]:
+                    return False
+                    
+                # Time for flux to decrease should be longer than that to increase
+                if end_idx - peak_idx <= peak_idx - start_idx:
+                    return False
+                    
+                # Check flux level consistency before and after flare
+                alter = (flux[peak_idx] - flux[start_idx - 2]) / (flux[peak_idx] - flux[end_idx + 2] + 1e-8)
+                # Flux level should be similar before and after flare
+                if alter < 0.5 or alter > 2 or np.isnan(alter):
+                    return False
+                    
+                # Check if the slope before peak is too steep
+                # if np.abs(flux[peak_idx] - flux[peak_idx-1]) < 1.2 * np.abs(flux[peak_idx-1] - flux[peak_idx-2]):
+                #     return False
+                    
+            except (IndexError, ValueError):
+                return False
+            
+            return True
+        
+        for i in range(predictions_np.shape[0]):
+            pred = predictions_np[i]
+            flux = flux_data[i, :, 0]  # Get flux data
+            flare_idx = np.where(pred == 1)[0]
+            
+            if len(flare_idx) == 0:
+                continue
+            
+            # Find continuous segments
+            splits = np.where(np.diff(flare_idx) > 1)[0] + 1
+            segments = np.split(flare_idx, splits)
+            
+            for segment in segments:
+                # Skip short segments early
+                if len(segment) < 3:
+                    continue
+                
+                start_idx = segment[0]
+                end_idx = segment[-1]
+                
+                # Find peak within segment
+                segment_flux = flux[start_idx:end_idx+1]
+                peak_idx = np.argmax(segment_flux) + start_idx
+                
+                # Validate flare characteristics
+                if not is_valid_flare(flux, start_idx, end_idx, peak_idx):
+                    continue
+                
+                # Valid flare event found
+                start_time = float(self.time_series[i][start_idx])
+                end_time = float(self.time_series[i][end_idx])
+                duration = end_time - start_time
+                event = {
+                    "start_idx": int(start_idx),
+                    "peak_idx": int(peak_idx),
+                    "end_idx": int(end_idx),
+                    "start_time": start_time,
+                    "peak_time": float(self.time_series[i][peak_idx]),
+                    "end_time": end_time,
+                    "duration": duration,
+                    "confidence": float(predictions[i, segment].mean()),
+                }
+                flare_events.append(event)
+        
+        return flare_events
+
+def load_flare_detection_pipeline(
+    model_name: str = "Maxwell-Jia/fcn4flare",
+    device: int = -1,
+    **kwargs
+) -> FlareDetectionPipeline:
+    """
+    Load a flare detection pipeline.
+    
+    Args:
+        model_name (str): The model name or path to load
+        device (int): Device to use (-1 for CPU, GPU number otherwise)
+        **kwargs: Additional arguments to pass to the pipeline
+        
+    Returns:
+        FlareDetectionPipeline: A pipeline for flare detection
+    """
+    return FlareDetectionPipeline(
+        model=model_name,
+        device=device,
+        **kwargs
+    )