Add source code and model

.gitignore

@@ -0,0 +1,70 @@
+AIGVDet/frame/
+AIGVDet/optical_result/
+AIGVDet/temp/
+frame/
+optical_result/
+temp/
+temp_uploads/
+
+miragenews/encodings/crops/
+miragenews/encodings/image/test1_nyt_mj/
+miragenews/encodings/image/test2_bbc_dalle/
+miragenews/encodings/image/test3_cnn_dalle/
+miragenews/encodings/image/test4_bbc_sdxl/
+miragenews/encodings/image/test5_cnn_sdxl/
+miragenews/encodings/image/train/
+miragenews/encodings/image/validation/
+miragenews/encodings/predictions/image/cbm-encoder/test1_nyt_mj/
+miragenews/encodings/predictions/image/cbm-encoder/test2_bbc_dalle/
+miragenews/encodings/predictions/image/cbm-encoder/test3_cnn_dalle/
+miragenews/encodings/predictions/image/cbm-encoder/test4_bbc_sdxl/
+miragenews/encodings/predictions/image/cbm-encoder/test5_cnn_sdxl/
+miragenews/encodings/predictions/image/cbm-encoder/train/
+miragenews/encodings/predictions/image/cbm-encoder/validation/
+miragenews/encodings/predictions/image/linear/validation/
+miragenews/encodings/predictions/image/linear/test1_nyt_mj/
+miragenews/encodings/predictions/image/linear/test2_bbc_dalle/
+miragenews/encodings/predictions/image/linear/test3_cnn_dalle/
+miragenews/encodings/predictions/image/linear/test4_bbc_sdxl/
+miragenews/encodings/predictions/image/linear/test5_cnn_sdxl/
+miragenews/encodings/predictions/image/linear/train/
+miragenews/encodings/predictions/image/merged/train/
+miragenews/encodings/predictions/image/merged/test1_nyt_mj/
+miragenews/encodings/predictions/image/merged/test2_bbc_dalle/
+miragenews/encodings/predictions/image/merged/test3_cnn_dalle/
+miragenews/encodings/predictions/image/merged/test4_bbc_sdxl/
+miragenews/encodings/predictions/image/merged/test5_cnn_sdxl/
+miragenews/encodings/predictions/image/merged/validation/
+
+miragenews/encodings/text/train/
+miragenews/encodings/text/test1_nyt_mj/
+miragenews/encodings/text/test2_bbc_dalle/
+miragenews/encodings/text/test3_cnn_dalle/
+miragenews/encodings/text/test4_bbc_sdxl/
+miragenews/encodings/text/test5_cnn_sdxl/
+miragenews/encodings/text/validation/
+
+miragenews/encodings/predictions/text/merged/train/
+miragenews/encodings/predictions/text/merged/test1_nyt_mj/
+miragenews/encodings/predictions/text/merged/test2_bbc_dalle/
+miragenews/encodings/predictions/text/merged/test3_cnn_dalle/
+miragenews/encodings/predictions/text/merged/test4_bbc_sdxl/
+miragenews/encodings/predictions/text/merged/test5_cnn_sdxl/
+miragenews/encodings/predictions/text/merged/validation/
+
+miragenews/encodings/predictions/text/tbm-encoder/train/
+miragenews/encodings/predictions/text/tbm-encoder/test1_nyt_mj/
+miragenews/encodings/predictions/text/tbm-encoder/test2_bbc_dalle/
+miragenews/encodings/predictions/text/tbm-encoder/test3_cnn_dalle/
+miragenews/encodings/predictions/text/tbm-encoder/test4_bbc_sdxl/
+miragenews/encodings/predictions/text/tbm-encoder/test5_cnn_sdxl/
+miragenews/encodings/predictions/text/tbm-encoder/validation/
+
+miragenews/encodings/predictions/text/linear/train/
+miragenews/encodings/predictions/text/linear/test1_nyt_mj/
+miragenews/encodings/predictions/text/linear/test2_bbc_dalle/
+miragenews/encodings/predictions/text/linear/test3_cnn_dalle/
+miragenews/encodings/predictions/text/linear/test4_bbc_sdxl/
+miragenews/encodings/predictions/text/linear/test5_cnn_sdxl/
+miragenews/encodings/predictions/text/linear/validation/
+AIGVDet/data/

AIGVDet/Dockerfile

@@ -0,0 +1,21 @@
+FROM pytorch/pytorch:2.2.0-cuda12.1-cudnn8-runtime
+
+# Install necessary OS packages for OpenCV
+RUN apt-get update && apt-get install -y \
+    libgl1 \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy all source code
+COPY . .
+
+# Default run command
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8003"]
+

AIGVDet/README.md

@@ -0,0 +1,94 @@
+## AIGVDet
+An official implementation code for paper "AI-Generated Video Detection via Spatial-Temporal Anomaly Learning", PRCV 2024. This repo will provide <B>codes, trained weights, and our training datasets</B>. 
+
+## Network Architecture
+<center> <img src="fig/NetworkArchitecture.png" alt="architecture"/> </center>
+
+## Dataset 
+- Download the preprocessed training frames from
+[Baiduyun Link](https://pan.baidu.com/s/17xmDyFjtcmNsoxmUeImMTQ?pwd=ra95) (extract code: ra95).
+- Download the test videos from [Google Drive](https://drive.google.com/drive/folders/1D84SRWEJ8BK8KBpTMuGi3BUM80mW_dKb?usp=sharing).
+
+**You are allowed to use the datasets for <B>research purpose only</B>.**
+
+## Training
+- Prepare for the training datasets.
+```
+└─data
+   ├── train
+   │   └── trainset_1
+   │       ├── 0_real
+   │       │   ├── video_00000
+   │       │   │    ├── 00000.png
+   │       │   │    └── ...
+   │       │   └── ...
+   │       └── 1_fake
+   │           ├── video_00000
+   │           │    ├── 00000.png
+   │           │    └── ...
+   │           └── ...
+   ├── val
+   │   └── val_set_1
+   │       ├── 0_real
+   │       │   ├── video_00000
+   │       │   │    ├── 00000.png
+   │       │   │    └── ...
+   │       │   └── ...
+   │       └── 1_fake
+   │           ├── video_00000
+   │           │    ├── 00000.png
+   │           │    └── ...
+   │           └── ...
+   └── test
+       └── testset_1
+           ├── 0_real
+           │   ├── video_00000
+           │   │    ├── 00000.png
+           │   │    └── ...
+           │   └── ...
+           └── 1_fake
+               ├── video_00000
+               │    ├── 00000.png
+               │    └── ...
+               └── ...
+
+```
+- Modify configuration file in `core/utils1/config.py`.
+- Train the Spatial Domain Detector with the RGB frames.
+```
+python train.py --gpus 0 --exp_name TRAIN_RGB_BRANCH datasets RGB_TRAINSET datasets_test RGB_TESTSET
+```
+- Train the Optical Flow Detector with the optical flow frames.
+```
+python train.py --gpus 0 --exp_name TRAIN_OF_BRANCH datasets OpticalFlow_TRAINSET datasets_test OpticalFlow_TESTSET
+```
+## Testing
+Download the weights from [Google Drive Link](https://drive.google.com/drive/folders/18JO_YxOEqwJYfbVvy308XjoV-N6fE4yP?usp=share_link) and move it into the `checkpoints/`.
+
+- Run on a dataset.
+Prepare the RGB frames and the optical flow maps.
+```
+python test.py -fop "data/test/hotshot" -mop "checkpoints/optical_aug.pth" -for "data/test/original/hotshot" -mor "checkpoints/original_aug.pth" -e "data/results/T2V/hotshot.csv" -ef "data/results/frame/T2V/hotshot.csv" -t 0.5
+```
+- Run on a video.
+Download the RAFT model weights from [Google Drive Link](https://drive.google.com/file/d/1MqDajR89k-xLV0HIrmJ0k-n8ZpG6_suM/view) and move it into the `raft_model/`.
+```
+python demo.py --use_cpu --path "video/000000.mp4" --folder_original_path "frame/000000" --folder_optical_flow_path "optical_result/000000" -mop "checkpoints/optical.pth" -mor "checkpoints/original.pth"
+```
+
+## License 
+The code and dataset is released only for academic research. Commercial usage is strictly prohibited.
+
+## Citation
+ ```
+@article{AIGVDet24,
+author = {Jianfa Bai and Man Lin and Gang Cao and Zijie Lou},
+title = {{AI-generated video detection via spatial-temporal anomaly learning}},
+conference = {The 7th Chinese Conference on Pattern Recognition and Computer Vision (PRCV)},
+year = {2024},}
+```
+
+## Contact
+If you have any questions, please contact us(lyan924@cuc.edu.cn).
+
+

AIGVDet/__init__.py

@@ -0,0 +1,3 @@
+from .main import run_video_to_json
+
+__all__ = ["run_video_to_json"]

AIGVDet/alt_cuda_corr/correlation.cpp

@@ -0,0 +1,54 @@
+#include <torch/extension.h>
+#include <vector>
+
+// CUDA forward declarations
+std::vector<torch::Tensor> corr_cuda_forward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    int radius);
+
+std::vector<torch::Tensor> corr_cuda_backward(
+  torch::Tensor fmap1,
+  torch::Tensor fmap2,
+  torch::Tensor coords,
+  torch::Tensor corr_grad,
+  int radius);
+
+// C++ interface
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> corr_forward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    int radius) {
+  CHECK_INPUT(fmap1);
+  CHECK_INPUT(fmap2);
+  CHECK_INPUT(coords);
+
+  return corr_cuda_forward(fmap1, fmap2, coords, radius);
+}
+
+
+std::vector<torch::Tensor> corr_backward(
+    torch::Tensor fmap1,
+    torch::Tensor fmap2,
+    torch::Tensor coords,
+    torch::Tensor corr_grad,
+    int radius) {
+  CHECK_INPUT(fmap1);
+  CHECK_INPUT(fmap2);
+  CHECK_INPUT(coords);
+  CHECK_INPUT(corr_grad);
+
+  return corr_cuda_backward(fmap1, fmap2, coords, corr_grad, radius);
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &corr_forward, "CORR forward");
+  m.def("backward", &corr_backward, "CORR backward");
+}

AIGVDet/alt_cuda_corr/correlation_kernel.cu

@@ -0,0 +1,324 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <vector>
+
+
+#define BLOCK_H 4
+#define BLOCK_W 8
+#define BLOCK_HW BLOCK_H * BLOCK_W
+#define CHANNEL_STRIDE 32
+
+
+__forceinline__ __device__
+bool within_bounds(int h, int w, int H, int W) {
+  return h >= 0 && h < H && w >= 0 && w < W;
+}
+
+template <typename scalar_t>
+__global__ void corr_forward_kernel(
+    const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap1,
+    const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap2,
+    const torch::PackedTensorAccessor32<scalar_t,5,torch::RestrictPtrTraits> coords,
+    torch::PackedTensorAccessor32<scalar_t,5,torch::RestrictPtrTraits> corr,
+    int r)
+{
+  const int b = blockIdx.x;
+  const int h0 = blockIdx.y * blockDim.x;
+  const int w0 = blockIdx.z * blockDim.y;
+  const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+  const int H1 = fmap1.size(1);
+  const int W1 = fmap1.size(2);
+  const int H2 = fmap2.size(1);
+  const int W2 = fmap2.size(2);
+  const int N = coords.size(1);
+  const int C = fmap1.size(3);
+
+  __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW+1];
+  __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW+1];
+  __shared__ scalar_t x2s[BLOCK_HW];
+  __shared__ scalar_t y2s[BLOCK_HW];
+
+  for (int c=0; c<C; c+=CHANNEL_STRIDE) {
+    for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+      int k1 = k + tid / CHANNEL_STRIDE;
+      int h1 = h0 + k1 / BLOCK_W;
+      int w1 = w0 + k1 % BLOCK_W;
+      int c1 = tid % CHANNEL_STRIDE;
+
+      auto fptr = fmap1[b][h1][w1];
+      if (within_bounds(h1, w1, H1, W1))
+        f1[c1][k1] = fptr[c+c1];
+      else
+        f1[c1][k1] = 0.0;
+    }
+
+    __syncthreads();
+
+    for (int n=0; n<N; n++) {
+      int h1 = h0 + threadIdx.x;
+      int w1 = w0 + threadIdx.y;
+      if (within_bounds(h1, w1, H1, W1)) {
+        x2s[tid] = coords[b][n][h1][w1][0];
+        y2s[tid] = coords[b][n][h1][w1][1];
+      }
+
+      scalar_t dx = x2s[tid] - floor(x2s[tid]);
+      scalar_t dy = y2s[tid] - floor(y2s[tid]);
+
+      int rd = 2*r + 1;
+      for (int iy=0; iy<rd+1; iy++) {
+        for (int ix=0; ix<rd+1; ix++) {
+          for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h2 = static_cast<int>(floor(y2s[k1]))-r+iy;
+            int w2 = static_cast<int>(floor(x2s[k1]))-r+ix;
+            int c2 = tid % CHANNEL_STRIDE;
+
+            auto fptr = fmap2[b][h2][w2];
+            if (within_bounds(h2, w2, H2, W2))
+              f2[c2][k1] = fptr[c+c2];
+            else
+              f2[c2][k1] = 0.0;
+          }
+
+          __syncthreads();
+      
+          scalar_t s = 0.0;
+          for (int k=0; k<CHANNEL_STRIDE; k++)
+            s += f1[k][tid] * f2[k][tid];
+
+          int ix_nw = H1*W1*((iy-1) + rd*(ix-1));
+          int ix_ne = H1*W1*((iy-1) + rd*ix);
+          int ix_sw = H1*W1*(iy + rd*(ix-1));
+          int ix_se = H1*W1*(iy + rd*ix);
+
+          scalar_t nw = s * (dy) * (dx);
+          scalar_t ne = s * (dy) * (1-dx);
+          scalar_t sw = s * (1-dy) * (dx);
+          scalar_t se = s * (1-dy) * (1-dx);
+
+          scalar_t* corr_ptr = &corr[b][n][0][h1][w1];
+
+          if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1))
+            *(corr_ptr + ix_nw) += nw;
+
+          if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1))
+            *(corr_ptr + ix_ne) += ne;
+
+          if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1))
+            *(corr_ptr + ix_sw) += sw;
+
+          if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+            *(corr_ptr + ix_se) += se;
+        }
+      } 
+    }
+  }
+}
+
+
+template <typename scalar_t>
+__global__ void corr_backward_kernel(
+    const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap1,
+    const torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap2,
+    const torch::PackedTensorAccessor32<scalar_t,5,torch::RestrictPtrTraits> coords,
+    const torch::PackedTensorAccessor32<scalar_t,5,torch::RestrictPtrTraits> corr_grad,
+    torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap1_grad,
+    torch::PackedTensorAccessor32<scalar_t,4,torch::RestrictPtrTraits> fmap2_grad,
+    torch::PackedTensorAccessor32<scalar_t,5,torch::RestrictPtrTraits> coords_grad,
+    int r)
+{
+
+  const int b = blockIdx.x;
+  const int h0 = blockIdx.y * blockDim.x;
+  const int w0 = blockIdx.z * blockDim.y;
+  const int tid = threadIdx.x * blockDim.y + threadIdx.y;
+
+  const int H1 = fmap1.size(1);
+  const int W1 = fmap1.size(2);
+  const int H2 = fmap2.size(1);
+  const int W2 = fmap2.size(2);
+  const int N = coords.size(1);
+  const int C = fmap1.size(3);
+
+  __shared__ scalar_t f1[CHANNEL_STRIDE][BLOCK_HW+1];
+  __shared__ scalar_t f2[CHANNEL_STRIDE][BLOCK_HW+1];
+
+  __shared__ scalar_t f1_grad[CHANNEL_STRIDE][BLOCK_HW+1];
+  __shared__ scalar_t f2_grad[CHANNEL_STRIDE][BLOCK_HW+1];
+
+  __shared__ scalar_t x2s[BLOCK_HW];
+  __shared__ scalar_t y2s[BLOCK_HW];
+
+  for (int c=0; c<C; c+=CHANNEL_STRIDE) {
+
+    for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+      int k1 = k + tid / CHANNEL_STRIDE;
+      int h1 = h0 + k1 / BLOCK_W;
+      int w1 = w0 + k1 % BLOCK_W;
+      int c1 = tid % CHANNEL_STRIDE;
+
+      auto fptr = fmap1[b][h1][w1];
+      if (within_bounds(h1, w1, H1, W1))
+        f1[c1][k1] = fptr[c+c1];
+      else
+        f1[c1][k1] = 0.0;
+
+      f1_grad[c1][k1] = 0.0;
+    }
+
+    __syncthreads();
+
+    int h1 = h0 + threadIdx.x;
+    int w1 = w0 + threadIdx.y;
+
+    for (int n=0; n<N; n++) {  
+      x2s[tid] = coords[b][n][h1][w1][0];
+      y2s[tid] = coords[b][n][h1][w1][1];
+
+      scalar_t dx = x2s[tid] - floor(x2s[tid]);
+      scalar_t dy = y2s[tid] - floor(y2s[tid]);
+
+      int rd = 2*r + 1;
+      for (int iy=0; iy<rd+1; iy++) {
+        for (int ix=0; ix<rd+1; ix++) {
+          for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h2 = static_cast<int>(floor(y2s[k1]))-r+iy;
+            int w2 = static_cast<int>(floor(x2s[k1]))-r+ix;
+            int c2 = tid % CHANNEL_STRIDE;
+
+            auto fptr = fmap2[b][h2][w2];
+            if (within_bounds(h2, w2, H2, W2))
+              f2[c2][k1] = fptr[c+c2];
+            else
+              f2[c2][k1] = 0.0;
+
+            f2_grad[c2][k1] = 0.0;
+          }
+
+          __syncthreads();
+      
+          const scalar_t* grad_ptr = &corr_grad[b][n][0][h1][w1];
+          scalar_t g = 0.0;
+
+          int ix_nw = H1*W1*((iy-1) + rd*(ix-1));
+          int ix_ne = H1*W1*((iy-1) + rd*ix);
+          int ix_sw = H1*W1*(iy + rd*(ix-1));
+          int ix_se = H1*W1*(iy + rd*ix);
+
+          if (iy > 0 && ix > 0 && within_bounds(h1, w1, H1, W1))
+            g +=  *(grad_ptr + ix_nw) * dy * dx;
+
+          if (iy > 0 && ix < rd && within_bounds(h1, w1, H1, W1))
+            g += *(grad_ptr + ix_ne) * dy * (1-dx);
+
+          if (iy < rd && ix > 0 && within_bounds(h1, w1, H1, W1))
+            g += *(grad_ptr + ix_sw) * (1-dy) * dx;
+
+          if (iy < rd && ix < rd && within_bounds(h1, w1, H1, W1))
+            g += *(grad_ptr + ix_se) * (1-dy) * (1-dx);
+            
+          for (int k=0; k<CHANNEL_STRIDE; k++) {
+            f1_grad[k][tid] += g * f2[k][tid];
+            f2_grad[k][tid] += g * f1[k][tid];
+          }
+
+          for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+            int k1 = k + tid / CHANNEL_STRIDE;
+            int h2 = static_cast<int>(floor(y2s[k1]))-r+iy;
+            int w2 = static_cast<int>(floor(x2s[k1]))-r+ix;
+            int c2 = tid % CHANNEL_STRIDE;
+
+            scalar_t* fptr = &fmap2_grad[b][h2][w2][0];
+            if (within_bounds(h2, w2, H2, W2))
+              atomicAdd(fptr+c+c2, f2_grad[c2][k1]);
+          }
+        }
+      } 
+    }
+    __syncthreads();
+
+
+    for (int k=0; k<BLOCK_HW; k+=BLOCK_HW/CHANNEL_STRIDE) {
+      int k1 = k + tid / CHANNEL_STRIDE;
+      int h1 = h0 + k1 / BLOCK_W;
+      int w1 = w0 + k1 % BLOCK_W;
+      int c1 = tid % CHANNEL_STRIDE;
+
+      scalar_t* fptr = &fmap1_grad[b][h1][w1][0];
+      if (within_bounds(h1, w1, H1, W1))
+        fptr[c+c1] += f1_grad[c1][k1];
+    }
+  }
+}
+
+
+
+std::vector<torch::Tensor> corr_cuda_forward(
+  torch::Tensor fmap1,
+  torch::Tensor fmap2,
+  torch::Tensor coords,
+  int radius)
+{
+  const auto B = coords.size(0);
+  const auto N = coords.size(1);
+  const auto H = coords.size(2);
+  const auto W = coords.size(3);
+
+  const auto rd = 2 * radius + 1;
+  auto opts = fmap1.options();
+  auto corr = torch::zeros({B, N, rd*rd, H, W}, opts);
+  
+  const dim3 blocks(B, (H+BLOCK_H-1)/BLOCK_H, (W+BLOCK_W-1)/BLOCK_W);
+  const dim3 threads(BLOCK_H, BLOCK_W);
+
+  corr_forward_kernel<float><<<blocks, threads>>>(
+    fmap1.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    fmap2.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    coords.packed_accessor32<float,5,torch::RestrictPtrTraits>(),
+    corr.packed_accessor32<float,5,torch::RestrictPtrTraits>(),
+    radius);
+
+  return {corr};
+}
+
+std::vector<torch::Tensor> corr_cuda_backward(
+  torch::Tensor fmap1,
+  torch::Tensor fmap2,
+  torch::Tensor coords,
+  torch::Tensor corr_grad,
+  int radius)
+{
+  const auto B = coords.size(0);
+  const auto N = coords.size(1);
+
+  const auto H1 = fmap1.size(1);
+  const auto W1 = fmap1.size(2);
+  const auto H2 = fmap2.size(1);
+  const auto W2 = fmap2.size(2);
+  const auto C = fmap1.size(3);
+
+  auto opts = fmap1.options();
+  auto fmap1_grad = torch::zeros({B, H1, W1, C}, opts);
+  auto fmap2_grad = torch::zeros({B, H2, W2, C}, opts);
+  auto coords_grad = torch::zeros({B, N, H1, W1, 2}, opts);
+    
+  const dim3 blocks(B, (H1+BLOCK_H-1)/BLOCK_H, (W1+BLOCK_W-1)/BLOCK_W);
+  const dim3 threads(BLOCK_H, BLOCK_W);
+
+
+  corr_backward_kernel<float><<<blocks, threads>>>(
+    fmap1.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    fmap2.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    coords.packed_accessor32<float,5,torch::RestrictPtrTraits>(),
+    corr_grad.packed_accessor32<float,5,torch::RestrictPtrTraits>(),
+    fmap1_grad.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    fmap2_grad.packed_accessor32<float,4,torch::RestrictPtrTraits>(),
+    coords_grad.packed_accessor32<float,5,torch::RestrictPtrTraits>(),
+    radius);
+
+  return {fmap1_grad, fmap2_grad, coords_grad};
+}

AIGVDet/alt_cuda_corr/setup.py

@@ -0,0 +1,15 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+setup(
+    name='correlation',
+    ext_modules=[
+        CUDAExtension('alt_cuda_corr',
+            sources=['correlation.cpp', 'correlation_kernel.cu'],
+            extra_compile_args={'cxx': [], 'nvcc': ['-O3']}),
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    })
+

AIGVDet/app.py

@@ -0,0 +1,155 @@
+# app.py
+import os
+import shutil
+import uuid
+import time
+from typing import Optional 
+
+from fastapi import FastAPI, UploadFile, File, HTTPException, status
+from pydantic import BaseModel
+from run import RUN
+
+
+class PredictionResponse(BaseModel):
+    authenticity_assessment: str
+    verification_tools_methods: str
+    synthetic_type: str
+    other_artifacts: str
+
+app = FastAPI(
+    title="Video Authenticity API",
+    description="Detect authentic vs synthetic video using deepfake detector.",
+)
+
+def predict(input_path: str) -> dict:
+    result_data = {
+        "authenticity_assessment": "Error",
+        "verification_tools_methods": "Deepfake Detection Model (Optical Flow + Frame Analysis)",
+        "synthetic_type": "N/A",
+        "other_artifacts": "Analysis failed."
+    }
+
+    if not os.path.isfile(input_path):
+        raise FileNotFoundError(f"File not found: {input_path}")
+
+    video_name = os.path.basename(input_path)
+    video_id = os.path.splitext(video_name)[0]
+
+    folder_original = f"frame/{video_id}"
+    folder_optical = f"optical_result/{video_id}"
+
+    args = [
+        "--path", input_path,
+        "--folder_original_path", folder_original,
+        "--folder_optical_flow_path", folder_optical,
+        "--model_optical_flow_path", "checkpoints/optical.pth",
+        "--model_original_path", "checkpoints/original.pth",
+    ]
+
+    try:
+        start_time = time.perf_counter()
+
+        output = RUN(args)
+
+        elapsed = time.perf_counter() - start_time
+        print(f"⏱️ [run_AIGVDetection] Service call took {elapsed:.2f} seconds")
+        
+        real_score = float(output.get("real_score", 0.0))
+        fake_score = float(output.get("fake_score", 0.0))
+
+        likely_authentic = real_score > fake_score
+        
+        if likely_authentic:
+            assessment = f"REAL (Authentic) | Confidence: {real_score:.4f}"
+            
+            analysis_text = (
+                "Our algorithms observed consistent and natural motion patterns across frames. "
+                "Inter-frame motion analysis indicates that objects maintain physical trajectories consistent with real-world recording, "
+                "without the jitter or warping artifacts typically associated with generative AI."
+            )
+            syn_type = "N/A"
+        else:
+            assessment = f"🤖 NOT REAL (Fake/Synthetic) | Confidence: {fake_score:.4f}"
+            
+            analysis_text = (
+                "Our algorithms have detected asynchronous and inconsistent movement between frames. "
+                "Upon conducting inter-frame motion analysis, we observed that objects and details within the video "
+                "fail to maintain natural motion trajectories. These anomalies—such as sudden velocity shifts, "
+                "subtle per-frame distortions, or motion vectors that defy physical laws—are characteristic indicators "
+                "typically found in AI-generated videos."
+            )
+            syn_type = "Video Deepfake / AI Generated"
+
+        tools = "Deepfake Detector (Optical Flow + CNN Frame Analysis)"
+
+        artifacts = (
+            f"{analysis_text}"
+        )
+
+        result_data = {
+            "authenticity_assessment": assessment,
+            "verification_tools_methods": tools,
+            "synthetic_type": syn_type,
+            "other_artifacts": artifacts
+        }
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        result_data["other_artifacts"] = f"Error during processing: {str(e)}"
+        
+    finally:
+        for folder in [folder_original, folder_optical]:
+            try:
+                if os.path.exists(folder):
+                    shutil.rmtree(folder)
+            except Exception as cleanup_error:
+                print(f"Error deleting folder {folder}: {cleanup_error}")
+
+    return result_data
+
+# --- API ENDPOINT ---
+@app.post("/predict", response_model=PredictionResponse)
+async def predict_endpoint(file: UploadFile = File(...)):
+    try:
+        for parent in ["frame", "optical_result", "uploads"]:
+            if os.path.exists(parent):
+                print(f"🧹 Cleaning folder: {parent}")
+                for item in os.listdir(parent):
+                    path = os.path.join(parent, item)
+                    try:
+                        if os.path.isfile(path) or os.path.islink(path):
+                            os.remove(path)
+                        elif os.path.isdir(path):
+                            shutil.rmtree(path)
+                    except Exception as e:
+                        print(f"⚠️ Cannot delete {path}: {e}")
+            else:
+                os.makedirs(parent)
+                print(f"📁 Created new folder: {parent}")
+
+        temp_filename = f"uploads_{uuid.uuid4().hex}_{file.filename}"
+        os.makedirs("uploads", exist_ok=True)
+        temp_filepath = os.path.join("uploads", temp_filename)
+
+        with open(temp_filepath, "wb") as buffer:
+            shutil.copyfileobj(file.file, buffer)
+
+        result = predict(temp_filepath)
+
+        if os.path.exists(temp_filepath):
+            os.remove(temp_filepath)
+
+        return result
+
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Prediction failed: {e}",
+        )
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=80022)

AIGVDet/checkpoints/optical.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23a167ba7adccb6421fd0eddd81fbfc69519d2a97a343dbf0d5da894b9893b19
+size 282581704

AIGVDet/checkpoints/original.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2df6f477a590f4b24b14eac9654b868a9f178312795df20749274a502d59bdd
+size 282581704

AIGVDet/core/corr.py

@@ -0,0 +1,91 @@
+import torch
+import torch.nn.functional as F
+from .utils.utils import bilinear_sampler, coords_grid
+
+try:
+    from .. import alt_cuda_corr
+except:
+    # alt_cuda_corr is not compiled
+    pass
+
+
+class CorrBlock:
+    def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
+        self.num_levels = num_levels
+        self.radius = radius
+        self.corr_pyramid = []
+
+        # all pairs correlation
+        corr = CorrBlock.corr(fmap1, fmap2)
+
+        batch, h1, w1, dim, h2, w2 = corr.shape
+        corr = corr.reshape(batch*h1*w1, dim, h2, w2)
+        
+        self.corr_pyramid.append(corr)
+        for i in range(self.num_levels-1):
+            corr = F.avg_pool2d(corr, 2, stride=2)
+            self.corr_pyramid.append(corr)
+
+    def __call__(self, coords):
+        r = self.radius
+        coords = coords.permute(0, 2, 3, 1)
+        batch, h1, w1, _ = coords.shape
+
+        out_pyramid = []
+        for i in range(self.num_levels):
+            corr = self.corr_pyramid[i]
+            dx = torch.linspace(-r, r, 2*r+1, device=coords.device)
+            dy = torch.linspace(-r, r, 2*r+1, device=coords.device)
+            delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
+
+            centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
+            delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
+            coords_lvl = centroid_lvl + delta_lvl
+
+            corr = bilinear_sampler(corr, coords_lvl)
+            corr = corr.view(batch, h1, w1, -1)
+            out_pyramid.append(corr)
+
+        out = torch.cat(out_pyramid, dim=-1)
+        return out.permute(0, 3, 1, 2).contiguous().float()
+
+    @staticmethod
+    def corr(fmap1, fmap2):
+        batch, dim, ht, wd = fmap1.shape
+        fmap1 = fmap1.view(batch, dim, ht*wd)
+        fmap2 = fmap2.view(batch, dim, ht*wd) 
+        
+        corr = torch.matmul(fmap1.transpose(1,2), fmap2)
+        corr = corr.view(batch, ht, wd, 1, ht, wd)
+        return corr  / torch.sqrt(torch.tensor(dim).float())
+
+
+class AlternateCorrBlock:
+    def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
+        self.num_levels = num_levels
+        self.radius = radius
+
+        self.pyramid = [(fmap1, fmap2)]
+        for i in range(self.num_levels):
+            fmap1 = F.avg_pool2d(fmap1, 2, stride=2)
+            fmap2 = F.avg_pool2d(fmap2, 2, stride=2)
+            self.pyramid.append((fmap1, fmap2))
+
+    def __call__(self, coords):
+        coords = coords.permute(0, 2, 3, 1)
+        B, H, W, _ = coords.shape
+        dim = self.pyramid[0][0].shape[1]
+
+        corr_list = []
+        for i in range(self.num_levels):
+            r = self.radius
+            fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous()
+            fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous()
+
+            coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous()
+            corr, = alt_cuda_corr.forward(fmap1_i, fmap2_i, coords_i, r)
+            corr_list.append(corr.squeeze(1))
+
+        corr = torch.stack(corr_list, dim=1)
+        corr = corr.reshape(B, -1, H, W)
+        return corr / torch.sqrt(torch.tensor(dim).float())

AIGVDet/core/datasets.py

@@ -0,0 +1,235 @@
+# Data loading based on https://github.com/NVIDIA/flownet2-pytorch
+
+import numpy as np
+import torch
+import torch.utils.data as data
+import torch.nn.functional as F
+
+import os
+import math
+import random
+from glob import glob
+import os.path as osp
+
+from utils import frame_utils
+from utils.augmentor import FlowAugmentor, SparseFlowAugmentor
+
+
+class FlowDataset(data.Dataset):
+    def __init__(self, aug_params=None, sparse=False):
+        self.augmentor = None
+        self.sparse = sparse
+        if aug_params is not None:
+            if sparse:
+                self.augmentor = SparseFlowAugmentor(**aug_params)
+            else:
+                self.augmentor = FlowAugmentor(**aug_params)
+
+        self.is_test = False
+        self.init_seed = False
+        self.flow_list = []
+        self.image_list = []
+        self.extra_info = []
+
+    def __getitem__(self, index):
+
+        if self.is_test:
+            img1 = frame_utils.read_gen(self.image_list[index][0])
+            img2 = frame_utils.read_gen(self.image_list[index][1])
+            img1 = np.array(img1).astype(np.uint8)[..., :3]
+            img2 = np.array(img2).astype(np.uint8)[..., :3]
+            img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
+            img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
+            return img1, img2, self.extra_info[index]
+
+        if not self.init_seed:
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                torch.manual_seed(worker_info.id)
+                np.random.seed(worker_info.id)
+                random.seed(worker_info.id)
+                self.init_seed = True
+
+        index = index % len(self.image_list)
+        valid = None
+        if self.sparse:
+            flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])
+        else:
+            flow = frame_utils.read_gen(self.flow_list[index])
+
+        img1 = frame_utils.read_gen(self.image_list[index][0])
+        img2 = frame_utils.read_gen(self.image_list[index][1])
+
+        flow = np.array(flow).astype(np.float32)
+        img1 = np.array(img1).astype(np.uint8)
+        img2 = np.array(img2).astype(np.uint8)
+
+        # grayscale images
+        if len(img1.shape) == 2:
+            img1 = np.tile(img1[...,None], (1, 1, 3))
+            img2 = np.tile(img2[...,None], (1, 1, 3))
+        else:
+            img1 = img1[..., :3]
+            img2 = img2[..., :3]
+
+        if self.augmentor is not None:
+            if self.sparse:
+                img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid)
+            else:
+                img1, img2, flow = self.augmentor(img1, img2, flow)
+
+        img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
+        img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
+        flow = torch.from_numpy(flow).permute(2, 0, 1).float()
+
+        if valid is not None:
+            valid = torch.from_numpy(valid)
+        else:
+            valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000)
+
+        return img1, img2, flow, valid.float()
+
+
+    def __rmul__(self, v):
+        self.flow_list = v * self.flow_list
+        self.image_list = v * self.image_list
+        return self
+        
+    def __len__(self):
+        return len(self.image_list)
+        
+
+class MpiSintel(FlowDataset):
+    def __init__(self, aug_params=None, split='training', root='datasets/Sintel', dstype='clean'):
+        super(MpiSintel, self).__init__(aug_params)
+        flow_root = osp.join(root, split, 'flow')
+        image_root = osp.join(root, split, dstype)
+
+        if split == 'test':
+            self.is_test = True
+
+        for scene in os.listdir(image_root):
+            image_list = sorted(glob(osp.join(image_root, scene, '*.png')))
+            for i in range(len(image_list)-1):
+                self.image_list += [ [image_list[i], image_list[i+1]] ]
+                self.extra_info += [ (scene, i) ] # scene and frame_id
+
+            if split != 'test':
+                self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))
+
+
+class FlyingChairs(FlowDataset):
+    def __init__(self, aug_params=None, split='train', root='datasets/FlyingChairs_release/data'):
+        super(FlyingChairs, self).__init__(aug_params)
+
+        images = sorted(glob(osp.join(root, '*.ppm')))
+        flows = sorted(glob(osp.join(root, '*.flo')))
+        assert (len(images)//2 == len(flows))
+
+        split_list = np.loadtxt('chairs_split.txt', dtype=np.int32)
+        for i in range(len(flows)):
+            xid = split_list[i]
+            if (split=='training' and xid==1) or (split=='validation' and xid==2):
+                self.flow_list += [ flows[i] ]
+                self.image_list += [ [images[2*i], images[2*i+1]] ]
+
+
+class FlyingThings3D(FlowDataset):
+    def __init__(self, aug_params=None, root='datasets/FlyingThings3D', dstype='frames_cleanpass'):
+        super(FlyingThings3D, self).__init__(aug_params)
+
+        for cam in ['left']:
+            for direction in ['into_future', 'into_past']:
+                image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*')))
+                image_dirs = sorted([osp.join(f, cam) for f in image_dirs])
+
+                flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))
+                flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])
+
+                for idir, fdir in zip(image_dirs, flow_dirs):
+                    images = sorted(glob(osp.join(idir, '*.png')) )
+                    flows = sorted(glob(osp.join(fdir, '*.pfm')) )
+                    for i in range(len(flows)-1):
+                        if direction == 'into_future':
+                            self.image_list += [ [images[i], images[i+1]] ]
+                            self.flow_list += [ flows[i] ]
+                        elif direction == 'into_past':
+                            self.image_list += [ [images[i+1], images[i]] ]
+                            self.flow_list += [ flows[i+1] ]
+      
+
+class KITTI(FlowDataset):
+    def __init__(self, aug_params=None, split='training', root='datasets/KITTI'):
+        super(KITTI, self).__init__(aug_params, sparse=True)
+        if split == 'testing':
+            self.is_test = True
+
+        root = osp.join(root, split)
+        images1 = sorted(glob(osp.join(root, 'image_2/*_10.png')))
+        images2 = sorted(glob(osp.join(root, 'image_2/*_11.png')))
+
+        for img1, img2 in zip(images1, images2):
+            frame_id = img1.split('/')[-1]
+            self.extra_info += [ [frame_id] ]
+            self.image_list += [ [img1, img2] ]
+
+        if split == 'training':
+            self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))
+
+
+class HD1K(FlowDataset):
+    def __init__(self, aug_params=None, root='datasets/HD1k'):
+        super(HD1K, self).__init__(aug_params, sparse=True)
+
+        seq_ix = 0
+        while 1:
+            flows = sorted(glob(os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix)))
+            images = sorted(glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix)))
+
+            if len(flows) == 0:
+                break
+
+            for i in range(len(flows)-1):
+                self.flow_list += [flows[i]]
+                self.image_list += [ [images[i], images[i+1]] ]
+
+            seq_ix += 1
+
+
+def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):
+    """ Create the data loader for the corresponding trainign set """
+
+    if args.stage == 'chairs':
+        aug_params = {'crop_size': args.image_size, 'min_scale': -0.1, 'max_scale': 1.0, 'do_flip': True}
+        train_dataset = FlyingChairs(aug_params, split='training')
+    
+    elif args.stage == 'things':
+        aug_params = {'crop_size': args.image_size, 'min_scale': -0.4, 'max_scale': 0.8, 'do_flip': True}
+        clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass')
+        final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass')
+        train_dataset = clean_dataset + final_dataset
+
+    elif args.stage == 'sintel':
+        aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.6, 'do_flip': True}
+        things = FlyingThings3D(aug_params, dstype='frames_cleanpass')
+        sintel_clean = MpiSintel(aug_params, split='training', dstype='clean')
+        sintel_final = MpiSintel(aug_params, split='training', dstype='final')        
+
+        if TRAIN_DS == 'C+T+K+S+H':
+            kitti = KITTI({'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.5, 'do_flip': True})
+            hd1k = HD1K({'crop_size': args.image_size, 'min_scale': -0.5, 'max_scale': 0.2, 'do_flip': True})
+            train_dataset = 100*sintel_clean + 100*sintel_final + 200*kitti + 5*hd1k + things
+
+        elif TRAIN_DS == 'C+T+K/S':
+            train_dataset = 100*sintel_clean + 100*sintel_final + things
+
+    elif args.stage == 'kitti':
+        aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.4, 'do_flip': False}
+        train_dataset = KITTI(aug_params, split='training')
+
+    train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size, 
+        pin_memory=False, shuffle=True, num_workers=4, drop_last=True)
+
+    print('Training with %d image pairs' % len(train_dataset))
+    return train_loader
+

AIGVDet/core/extractor.py

@@ -0,0 +1,267 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(ResidualBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes)
+            self.norm2 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes)
+            self.norm2 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm3 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            if not stride == 1:
+                self.norm3 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+
+
+class BottleneckBlock(nn.Module):
+    def __init__(self, in_planes, planes, norm_fn='group', stride=1):
+        super(BottleneckBlock, self).__init__()
+  
+        self.conv1 = nn.Conv2d(in_planes, planes//4, kernel_size=1, padding=0)
+        self.conv2 = nn.Conv2d(planes//4, planes//4, kernel_size=3, padding=1, stride=stride)
+        self.conv3 = nn.Conv2d(planes//4, planes, kernel_size=1, padding=0)
+        self.relu = nn.ReLU(inplace=True)
+
+        num_groups = planes // 8
+
+        if norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes//4)
+            self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+            if not stride == 1:
+                self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+        
+        elif norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(planes//4)
+            self.norm2 = nn.BatchNorm2d(planes//4)
+            self.norm3 = nn.BatchNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.BatchNorm2d(planes)
+        
+        elif norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(planes//4)
+            self.norm2 = nn.InstanceNorm2d(planes//4)
+            self.norm3 = nn.InstanceNorm2d(planes)
+            if not stride == 1:
+                self.norm4 = nn.InstanceNorm2d(planes)
+
+        elif norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+            self.norm2 = nn.Sequential()
+            self.norm3 = nn.Sequential()
+            if not stride == 1:
+                self.norm4 = nn.Sequential()
+
+        if stride == 1:
+            self.downsample = None
+        
+        else:    
+            self.downsample = nn.Sequential(
+                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4)
+
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.norm1(self.conv1(y)))
+        y = self.relu(self.norm2(self.conv2(y)))
+        y = self.relu(self.norm3(self.conv3(y)))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x+y)
+
+class BasicEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
+        super(BasicEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(64)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(64)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 64
+        self.layer1 = self._make_layer(64,  stride=1)
+        self.layer2 = self._make_layer(96, stride=2)
+        self.layer3 = self._make_layer(128, stride=2)
+
+        # output convolution
+        self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+        
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+
+        # if input is list, combine batch dimension
+        is_list = isinstance(x, tuple) or isinstance(x, list)
+        if is_list:
+            batch_dim = x[0].shape[0]
+            x = torch.cat(x, dim=0)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+
+        x = self.conv2(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+
+        if is_list:
+            x = torch.split(x, [batch_dim, batch_dim], dim=0)
+
+        return x
+
+
+class SmallEncoder(nn.Module):
+    def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0):
+        super(SmallEncoder, self).__init__()
+        self.norm_fn = norm_fn
+
+        if self.norm_fn == 'group':
+            self.norm1 = nn.GroupNorm(num_groups=8, num_channels=32)
+            
+        elif self.norm_fn == 'batch':
+            self.norm1 = nn.BatchNorm2d(32)
+
+        elif self.norm_fn == 'instance':
+            self.norm1 = nn.InstanceNorm2d(32)
+
+        elif self.norm_fn == 'none':
+            self.norm1 = nn.Sequential()
+
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=7, stride=2, padding=3)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.in_planes = 32
+        self.layer1 = self._make_layer(32,  stride=1)
+        self.layer2 = self._make_layer(64, stride=2)
+        self.layer3 = self._make_layer(96, stride=2)
+
+        self.dropout = None
+        if dropout > 0:
+            self.dropout = nn.Dropout2d(p=dropout)
+        
+        self.conv2 = nn.Conv2d(96, output_dim, kernel_size=1)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)):
+                if m.weight is not None:
+                    nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, dim, stride=1):
+        layer1 = BottleneckBlock(self.in_planes, dim, self.norm_fn, stride=stride)
+        layer2 = BottleneckBlock(dim, dim, self.norm_fn, stride=1)
+        layers = (layer1, layer2)
+    
+        self.in_planes = dim
+        return nn.Sequential(*layers)
+
+
+    def forward(self, x):
+
+        # if input is list, combine batch dimension
+        is_list = isinstance(x, tuple) or isinstance(x, list)
+        if is_list:
+            batch_dim = x[0].shape[0]
+            x = torch.cat(x, dim=0)
+
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.relu1(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.conv2(x)
+
+        if self.training and self.dropout is not None:
+            x = self.dropout(x)
+
+        if is_list:
+            x = torch.split(x, [batch_dim, batch_dim], dim=0)
+
+        return x

AIGVDet/core/raft.py

@@ -0,0 +1,144 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .update import BasicUpdateBlock, SmallUpdateBlock
+from .extractor import BasicEncoder, SmallEncoder
+from .corr import CorrBlock, AlternateCorrBlock
+from .utils.utils import bilinear_sampler, coords_grid, upflow8
+
+try:
+    autocast = torch.cuda.amp.autocast
+except:
+    # dummy autocast for PyTorch < 1.6
+    class autocast:
+        def __init__(self, enabled):
+            pass
+        def __enter__(self):
+            pass
+        def __exit__(self, *args):
+            pass
+
+
+class RAFT(nn.Module):
+    def __init__(self, args):
+        super(RAFT, self).__init__()
+        self.args = args
+
+        if args.small:
+            self.hidden_dim = hdim = 96
+            self.context_dim = cdim = 64
+            args.corr_levels = 4
+            args.corr_radius = 3
+        
+        else:
+            self.hidden_dim = hdim = 128
+            self.context_dim = cdim = 128
+            args.corr_levels = 4
+            args.corr_radius = 4
+
+        if 'dropout' not in self.args:
+            self.args.dropout = 0
+
+        if 'alternate_corr' not in self.args:
+            self.args.alternate_corr = False
+
+        # feature network, context network, and update block
+        if args.small:
+            self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout)        
+            self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout)
+            self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)
+
+        else:
+            self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout)        
+            self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
+            self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)
+
+    def freeze_bn(self):
+        for m in self.modules():
+            if isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def initialize_flow(self, img):
+        """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
+        N, C, H, W = img.shape
+        coords0 = coords_grid(N, H//8, W//8, device=img.device)
+        coords1 = coords_grid(N, H//8, W//8, device=img.device)
+
+        # optical flow computed as difference: flow = coords1 - coords0
+        return coords0, coords1
+
+    def upsample_flow(self, flow, mask):
+        """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
+        N, _, H, W = flow.shape
+        mask = mask.view(N, 1, 9, 8, 8, H, W)
+        mask = torch.softmax(mask, dim=2)
+
+        up_flow = F.unfold(8 * flow, [3,3], padding=1)
+        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)
+
+        up_flow = torch.sum(mask * up_flow, dim=2)
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
+        return up_flow.reshape(N, 2, 8*H, 8*W)
+
+
+    def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
+        """ Estimate optical flow between pair of frames """
+
+        image1 = 2 * (image1 / 255.0) - 1.0
+        image2 = 2 * (image2 / 255.0) - 1.0
+
+        image1 = image1.contiguous()
+        image2 = image2.contiguous()
+
+        hdim = self.hidden_dim
+        cdim = self.context_dim
+
+        # run the feature network
+        with autocast(enabled=self.args.mixed_precision):
+            fmap1, fmap2 = self.fnet([image1, image2])        
+        
+        fmap1 = fmap1.float()
+        fmap2 = fmap2.float()
+        if self.args.alternate_corr:
+            corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
+        else:
+            corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
+
+        # run the context network
+        with autocast(enabled=self.args.mixed_precision):
+            cnet = self.cnet(image1)
+            net, inp = torch.split(cnet, [hdim, cdim], dim=1)
+            net = torch.tanh(net)
+            inp = torch.relu(inp)
+
+        coords0, coords1 = self.initialize_flow(image1)
+
+        if flow_init is not None:
+            coords1 = coords1 + flow_init
+
+        flow_predictions = []
+        for itr in range(iters):
+            coords1 = coords1.detach()
+            corr = corr_fn(coords1) # index correlation volume
+
+            flow = coords1 - coords0
+            with autocast(enabled=self.args.mixed_precision):
+                net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
+
+            # F(t+1) = F(t) + \Delta(t)
+            coords1 = coords1 + delta_flow
+
+            # upsample predictions
+            if up_mask is None:
+                flow_up = upflow8(coords1 - coords0)
+            else:
+                flow_up = self.upsample_flow(coords1 - coords0, up_mask)
+            
+            flow_predictions.append(flow_up)
+
+        if test_mode:
+            return coords1 - coords0, flow_up
+            
+        return flow_predictions

AIGVDet/core/update.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class FlowHead(nn.Module):
+    def __init__(self, input_dim=128, hidden_dim=256):
+        super(FlowHead, self).__init__()
+        self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
+        self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.conv2(self.relu(self.conv1(x)))
+
+class ConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=192+128):
+        super(ConvGRU, self).__init__()
+        self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+        self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+        self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, 3, padding=1)
+
+    def forward(self, h, x):
+        hx = torch.cat([h, x], dim=1)
+
+        z = torch.sigmoid(self.convz(hx))
+        r = torch.sigmoid(self.convr(hx))
+        q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
+
+        h = (1-z) * h + z * q
+        return h
+
+class SepConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=192+128):
+        super(SepConvGRU, self).__init__()
+        self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+        self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+        self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
+
+        self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+        self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+        self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
+
+
+    def forward(self, h, x):
+        # horizontal
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz1(hx))
+        r = torch.sigmoid(self.convr1(hx))
+        q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))        
+        h = (1-z) * h + z * q
+
+        # vertical
+        hx = torch.cat([h, x], dim=1)
+        z = torch.sigmoid(self.convz2(hx))
+        r = torch.sigmoid(self.convr2(hx))
+        q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))       
+        h = (1-z) * h + z * q
+
+        return h
+
+class SmallMotionEncoder(nn.Module):
+    def __init__(self, args):
+        super(SmallMotionEncoder, self).__init__()
+        cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
+        self.convc1 = nn.Conv2d(cor_planes, 96, 1, padding=0)
+        self.convf1 = nn.Conv2d(2, 64, 7, padding=3)
+        self.convf2 = nn.Conv2d(64, 32, 3, padding=1)
+        self.conv = nn.Conv2d(128, 80, 3, padding=1)
+
+    def forward(self, flow, corr):
+        cor = F.relu(self.convc1(corr))
+        flo = F.relu(self.convf1(flow))
+        flo = F.relu(self.convf2(flo))
+        cor_flo = torch.cat([cor, flo], dim=1)
+        out = F.relu(self.conv(cor_flo))
+        return torch.cat([out, flow], dim=1)
+
+class BasicMotionEncoder(nn.Module):
+    def __init__(self, args):
+        super(BasicMotionEncoder, self).__init__()
+        cor_planes = args.corr_levels * (2*args.corr_radius + 1)**2
+        self.convc1 = nn.Conv2d(cor_planes, 256, 1, padding=0)
+        self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
+        self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
+        self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv = nn.Conv2d(64+192, 128-2, 3, padding=1)
+
+    def forward(self, flow, corr):
+        cor = F.relu(self.convc1(corr))
+        cor = F.relu(self.convc2(cor))
+        flo = F.relu(self.convf1(flow))
+        flo = F.relu(self.convf2(flo))
+
+        cor_flo = torch.cat([cor, flo], dim=1)
+        out = F.relu(self.conv(cor_flo))
+        return torch.cat([out, flow], dim=1)
+
+class SmallUpdateBlock(nn.Module):
+    def __init__(self, args, hidden_dim=96):
+        super(SmallUpdateBlock, self).__init__()
+        self.encoder = SmallMotionEncoder(args)
+        self.gru = ConvGRU(hidden_dim=hidden_dim, input_dim=82+64)
+        self.flow_head = FlowHead(hidden_dim, hidden_dim=128)
+
+    def forward(self, net, inp, corr, flow):
+        motion_features = self.encoder(flow, corr)
+        inp = torch.cat([inp, motion_features], dim=1)
+        net = self.gru(net, inp)
+        delta_flow = self.flow_head(net)
+
+        return net, None, delta_flow
+
+class BasicUpdateBlock(nn.Module):
+    def __init__(self, args, hidden_dim=128, input_dim=128):
+        super(BasicUpdateBlock, self).__init__()
+        self.args = args
+        self.encoder = BasicMotionEncoder(args)
+        self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
+        self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
+
+        self.mask = nn.Sequential(
+            nn.Conv2d(128, 256, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 64*9, 1, padding=0))
+
+    def forward(self, net, inp, corr, flow, upsample=True):
+        motion_features = self.encoder(flow, corr)
+        inp = torch.cat([inp, motion_features], dim=1)
+
+        net = self.gru(net, inp)
+        delta_flow = self.flow_head(net)
+
+        # scale mask to balence gradients
+        mask = .25 * self.mask(net)
+        return net, mask, delta_flow
+
+
+

AIGVDet/core/utils/augmentor.py

@@ -0,0 +1,246 @@
+import numpy as np
+import random
+import math
+from PIL import Image
+
+import cv2
+cv2.setNumThreads(0)
+cv2.ocl.setUseOpenCL(False)
+
+import torch
+from torchvision.transforms import ColorJitter
+import torch.nn.functional as F
+
+
+class FlowAugmentor:
+    def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=True):
+        
+        # spatial augmentation params
+        self.crop_size = crop_size
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.spatial_aug_prob = 0.8
+        self.stretch_prob = 0.8
+        self.max_stretch = 0.2
+
+        # flip augmentation params
+        self.do_flip = do_flip
+        self.h_flip_prob = 0.5
+        self.v_flip_prob = 0.1
+
+        # photometric augmentation params
+        self.photo_aug = ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5/3.14)
+        self.asymmetric_color_aug_prob = 0.2
+        self.eraser_aug_prob = 0.5
+
+    def color_transform(self, img1, img2):
+        """ Photometric augmentation """
+
+        # asymmetric
+        if np.random.rand() < self.asymmetric_color_aug_prob:
+            img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
+            img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)
+
+        # symmetric
+        else:
+            image_stack = np.concatenate([img1, img2], axis=0)
+            image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
+            img1, img2 = np.split(image_stack, 2, axis=0)
+
+        return img1, img2
+
+    def eraser_transform(self, img1, img2, bounds=[50, 100]):
+        """ Occlusion augmentation """
+
+        ht, wd = img1.shape[:2]
+        if np.random.rand() < self.eraser_aug_prob:
+            mean_color = np.mean(img2.reshape(-1, 3), axis=0)
+            for _ in range(np.random.randint(1, 3)):
+                x0 = np.random.randint(0, wd)
+                y0 = np.random.randint(0, ht)
+                dx = np.random.randint(bounds[0], bounds[1])
+                dy = np.random.randint(bounds[0], bounds[1])
+                img2[y0:y0+dy, x0:x0+dx, :] = mean_color
+
+        return img1, img2
+
+    def spatial_transform(self, img1, img2, flow):
+        # randomly sample scale
+        ht, wd = img1.shape[:2]
+        min_scale = np.maximum(
+            (self.crop_size[0] + 8) / float(ht), 
+            (self.crop_size[1] + 8) / float(wd))
+
+        scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
+        scale_x = scale
+        scale_y = scale
+        if np.random.rand() < self.stretch_prob:
+            scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
+            scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
+        
+        scale_x = np.clip(scale_x, min_scale, None)
+        scale_y = np.clip(scale_y, min_scale, None)
+
+        if np.random.rand() < self.spatial_aug_prob:
+            # rescale the images
+            img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+            img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+            flow = cv2.resize(flow, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+            flow = flow * [scale_x, scale_y]
+
+        if self.do_flip:
+            if np.random.rand() < self.h_flip_prob: # h-flip
+                img1 = img1[:, ::-1]
+                img2 = img2[:, ::-1]
+                flow = flow[:, ::-1] * [-1.0, 1.0]
+
+            if np.random.rand() < self.v_flip_prob: # v-flip
+                img1 = img1[::-1, :]
+                img2 = img2[::-1, :]
+                flow = flow[::-1, :] * [1.0, -1.0]
+
+        y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
+        x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
+        
+        img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+
+        return img1, img2, flow
+
+    def __call__(self, img1, img2, flow):
+        img1, img2 = self.color_transform(img1, img2)
+        img1, img2 = self.eraser_transform(img1, img2)
+        img1, img2, flow = self.spatial_transform(img1, img2, flow)
+
+        img1 = np.ascontiguousarray(img1)
+        img2 = np.ascontiguousarray(img2)
+        flow = np.ascontiguousarray(flow)
+
+        return img1, img2, flow
+
+class SparseFlowAugmentor:
+    def __init__(self, crop_size, min_scale=-0.2, max_scale=0.5, do_flip=False):
+        # spatial augmentation params
+        self.crop_size = crop_size
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.spatial_aug_prob = 0.8
+        self.stretch_prob = 0.8
+        self.max_stretch = 0.2
+
+        # flip augmentation params
+        self.do_flip = do_flip
+        self.h_flip_prob = 0.5
+        self.v_flip_prob = 0.1
+
+        # photometric augmentation params
+        self.photo_aug = ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3/3.14)
+        self.asymmetric_color_aug_prob = 0.2
+        self.eraser_aug_prob = 0.5
+        
+    def color_transform(self, img1, img2):
+        image_stack = np.concatenate([img1, img2], axis=0)
+        image_stack = np.array(self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8)
+        img1, img2 = np.split(image_stack, 2, axis=0)
+        return img1, img2
+
+    def eraser_transform(self, img1, img2):
+        ht, wd = img1.shape[:2]
+        if np.random.rand() < self.eraser_aug_prob:
+            mean_color = np.mean(img2.reshape(-1, 3), axis=0)
+            for _ in range(np.random.randint(1, 3)):
+                x0 = np.random.randint(0, wd)
+                y0 = np.random.randint(0, ht)
+                dx = np.random.randint(50, 100)
+                dy = np.random.randint(50, 100)
+                img2[y0:y0+dy, x0:x0+dx, :] = mean_color
+
+        return img1, img2
+
+    def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0):
+        ht, wd = flow.shape[:2]
+        coords = np.meshgrid(np.arange(wd), np.arange(ht))
+        coords = np.stack(coords, axis=-1)
+
+        coords = coords.reshape(-1, 2).astype(np.float32)
+        flow = flow.reshape(-1, 2).astype(np.float32)
+        valid = valid.reshape(-1).astype(np.float32)
+
+        coords0 = coords[valid>=1]
+        flow0 = flow[valid>=1]
+
+        ht1 = int(round(ht * fy))
+        wd1 = int(round(wd * fx))
+
+        coords1 = coords0 * [fx, fy]
+        flow1 = flow0 * [fx, fy]
+
+        xx = np.round(coords1[:,0]).astype(np.int32)
+        yy = np.round(coords1[:,1]).astype(np.int32)
+
+        v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
+        xx = xx[v]
+        yy = yy[v]
+        flow1 = flow1[v]
+
+        flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32)
+        valid_img = np.zeros([ht1, wd1], dtype=np.int32)
+
+        flow_img[yy, xx] = flow1
+        valid_img[yy, xx] = 1
+
+        return flow_img, valid_img
+
+    def spatial_transform(self, img1, img2, flow, valid):
+        # randomly sample scale
+
+        ht, wd = img1.shape[:2]
+        min_scale = np.maximum(
+            (self.crop_size[0] + 1) / float(ht), 
+            (self.crop_size[1] + 1) / float(wd))
+
+        scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
+        scale_x = np.clip(scale, min_scale, None)
+        scale_y = np.clip(scale, min_scale, None)
+
+        if np.random.rand() < self.spatial_aug_prob:
+            # rescale the images
+            img1 = cv2.resize(img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+            img2 = cv2.resize(img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR)
+            flow, valid = self.resize_sparse_flow_map(flow, valid, fx=scale_x, fy=scale_y)
+
+        if self.do_flip:
+            if np.random.rand() < 0.5: # h-flip
+                img1 = img1[:, ::-1]
+                img2 = img2[:, ::-1]
+                flow = flow[:, ::-1] * [-1.0, 1.0]
+                valid = valid[:, ::-1]
+
+        margin_y = 20
+        margin_x = 50
+
+        y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0] + margin_y)
+        x0 = np.random.randint(-margin_x, img1.shape[1] - self.crop_size[1] + margin_x)
+
+        y0 = np.clip(y0, 0, img1.shape[0] - self.crop_size[0])
+        x0 = np.clip(x0, 0, img1.shape[1] - self.crop_size[1])
+
+        img1 = img1[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        img2 = img2[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        flow = flow[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        valid = valid[y0:y0+self.crop_size[0], x0:x0+self.crop_size[1]]
+        return img1, img2, flow, valid
+
+
+    def __call__(self, img1, img2, flow, valid):
+        img1, img2 = self.color_transform(img1, img2)
+        img1, img2 = self.eraser_transform(img1, img2)
+        img1, img2, flow, valid = self.spatial_transform(img1, img2, flow, valid)
+
+        img1 = np.ascontiguousarray(img1)
+        img2 = np.ascontiguousarray(img2)
+        flow = np.ascontiguousarray(flow)
+        valid = np.ascontiguousarray(valid)
+
+        return img1, img2, flow, valid

AIGVDet/core/utils/flow_viz.py

@@ -0,0 +1,132 @@
+# Flow visualization code used from https://github.com/tomrunia/OpticalFlow_Visualization
+
+
+# MIT License
+#
+# Copyright (c) 2018 Tom Runia
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to conditions.
+#
+# Author: Tom Runia
+# Date Created: 2018-08-03
+
+import numpy as np
+
+def make_colorwheel():
+    """
+    Generates a color wheel for optical flow visualization as presented in:
+        Baker et al. "A Database and Evaluation Methodology for Optical Flow" (ICCV, 2007)
+        URL: http://vision.middlebury.edu/flow/flowEval-iccv07.pdf
+
+    Code follows the original C++ source code of Daniel Scharstein.
+    Code follows the the Matlab source code of Deqing Sun.
+
+    Returns:
+        np.ndarray: Color wheel
+    """
+
+    RY = 15
+    YG = 6
+    GC = 4
+    CB = 11
+    BM = 13
+    MR = 6
+
+    ncols = RY + YG + GC + CB + BM + MR
+    colorwheel = np.zeros((ncols, 3))
+    col = 0
+
+    # RY
+    colorwheel[0:RY, 0] = 255
+    colorwheel[0:RY, 1] = np.floor(255*np.arange(0,RY)/RY)
+    col = col+RY
+    # YG
+    colorwheel[col:col+YG, 0] = 255 - np.floor(255*np.arange(0,YG)/YG)
+    colorwheel[col:col+YG, 1] = 255
+    col = col+YG
+    # GC
+    colorwheel[col:col+GC, 1] = 255
+    colorwheel[col:col+GC, 2] = np.floor(255*np.arange(0,GC)/GC)
+    col = col+GC
+    # CB
+    colorwheel[col:col+CB, 1] = 255 - np.floor(255*np.arange(CB)/CB)
+    colorwheel[col:col+CB, 2] = 255
+    col = col+CB
+    # BM
+    colorwheel[col:col+BM, 2] = 255
+    colorwheel[col:col+BM, 0] = np.floor(255*np.arange(0,BM)/BM)
+    col = col+BM
+    # MR
+    colorwheel[col:col+MR, 2] = 255 - np.floor(255*np.arange(MR)/MR)
+    colorwheel[col:col+MR, 0] = 255
+    return colorwheel
+
+
+def flow_uv_to_colors(u, v, convert_to_bgr=False):
+    """
+    Applies the flow color wheel to (possibly clipped) flow components u and v.
+
+    According to the C++ source code of Daniel Scharstein
+    According to the Matlab source code of Deqing Sun
+
+    Args:
+        u (np.ndarray): Input horizontal flow of shape [H,W]
+        v (np.ndarray): Input vertical flow of shape [H,W]
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    flow_image = np.zeros((u.shape[0], u.shape[1], 3), np.uint8)
+    colorwheel = make_colorwheel()  # shape [55x3]
+    ncols = colorwheel.shape[0]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    a = np.arctan2(-v, -u)/np.pi
+    fk = (a+1) / 2*(ncols-1)
+    k0 = np.floor(fk).astype(np.int32)
+    k1 = k0 + 1
+    k1[k1 == ncols] = 0
+    f = fk - k0
+    for i in range(colorwheel.shape[1]):
+        tmp = colorwheel[:,i]
+        col0 = tmp[k0] / 255.0
+        col1 = tmp[k1] / 255.0
+        col = (1-f)*col0 + f*col1
+        idx = (rad <= 1)
+        col[idx]  = 1 - rad[idx] * (1-col[idx])
+        col[~idx] = col[~idx] * 0.75   # out of range
+        # Note the 2-i => BGR instead of RGB
+        ch_idx = 2-i if convert_to_bgr else i
+        flow_image[:,:,ch_idx] = np.floor(255 * col)
+    return flow_image
+
+
+def flow_to_image(flow_uv, clip_flow=None, convert_to_bgr=False):
+    """
+    Expects a two dimensional flow image of shape.
+
+    Args:
+        flow_uv (np.ndarray): Flow UV image of shape [H,W,2]
+        clip_flow (float, optional): Clip maximum of flow values. Defaults to None.
+        convert_to_bgr (bool, optional): Convert output image to BGR. Defaults to False.
+
+    Returns:
+        np.ndarray: Flow visualization image of shape [H,W,3]
+    """
+    assert flow_uv.ndim == 3, 'input flow must have three dimensions'
+    assert flow_uv.shape[2] == 2, 'input flow must have shape [H,W,2]'
+    if clip_flow is not None:
+        flow_uv = np.clip(flow_uv, 0, clip_flow)
+    u = flow_uv[:,:,0]
+    v = flow_uv[:,:,1]
+    rad = np.sqrt(np.square(u) + np.square(v))
+    rad_max = np.max(rad)
+    epsilon = 1e-5
+    u = u / (rad_max + epsilon)
+    v = v / (rad_max + epsilon)
+    return flow_uv_to_colors(u, v, convert_to_bgr)

AIGVDet/core/utils/frame_utils.py

@@ -0,0 +1,137 @@
+import numpy as np
+from PIL import Image
+from os.path import *
+import re
+
+import cv2
+cv2.setNumThreads(0)
+cv2.ocl.setUseOpenCL(False)
+
+TAG_CHAR = np.array([202021.25], np.float32)
+
+def readFlow(fn):
+    """ Read .flo file in Middlebury format"""
+    # Code adapted from:
+    # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
+
+    # WARNING: this will work on little-endian architectures (eg Intel x86) only!
+    # print 'fn = %s'%(fn)
+    with open(fn, 'rb') as f:
+        magic = np.fromfile(f, np.float32, count=1)
+        if 202021.25 != magic:
+            print('Magic number incorrect. Invalid .flo file')
+            return None
+        else:
+            w = np.fromfile(f, np.int32, count=1)
+            h = np.fromfile(f, np.int32, count=1)
+            # print 'Reading %d x %d flo file\n' % (w, h)
+            data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
+            # Reshape data into 3D array (columns, rows, bands)
+            # The reshape here is for visualization, the original code is (w,h,2)
+            return np.resize(data, (int(h), int(w), 2))
+
+def readPFM(file):
+    file = open(file, 'rb')
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header == b'PF':
+        color = True
+    elif header == b'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
+    if dim_match:
+        width, height = map(int, dim_match.groups())
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().rstrip())
+    if scale < 0: # little-endian
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>' # big-endian
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data
+
+def writeFlow(filename,uv,v=None):
+    """ Write optical flow to file.
+    
+    If v is None, uv is assumed to contain both u and v channels,
+    stacked in depth.
+    Original code by Deqing Sun, adapted from Daniel Scharstein.
+    """
+    nBands = 2
+
+    if v is None:
+        assert(uv.ndim == 3)
+        assert(uv.shape[2] == 2)
+        u = uv[:,:,0]
+        v = uv[:,:,1]
+    else:
+        u = uv
+
+    assert(u.shape == v.shape)
+    height,width = u.shape
+    f = open(filename,'wb')
+    # write the header
+    f.write(TAG_CHAR)
+    np.array(width).astype(np.int32).tofile(f)
+    np.array(height).astype(np.int32).tofile(f)
+    # arrange into matrix form
+    tmp = np.zeros((height, width*nBands))
+    tmp[:,np.arange(width)*2] = u
+    tmp[:,np.arange(width)*2 + 1] = v
+    tmp.astype(np.float32).tofile(f)
+    f.close()
+
+
+def readFlowKITTI(filename):
+    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
+    flow = flow[:,:,::-1].astype(np.float32)
+    flow, valid = flow[:, :, :2], flow[:, :, 2]
+    flow = (flow - 2**15) / 64.0
+    return flow, valid
+
+def readDispKITTI(filename):
+    disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
+    valid = disp > 0.0
+    flow = np.stack([-disp, np.zeros_like(disp)], -1)
+    return flow, valid
+
+
+def writeFlowKITTI(filename, uv):
+    uv = 64.0 * uv + 2**15
+    valid = np.ones([uv.shape[0], uv.shape[1], 1])
+    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
+    cv2.imwrite(filename, uv[..., ::-1])
+    
+
+def read_gen(file_name, pil=False):
+    ext = splitext(file_name)[-1]
+    if ext == '.png' or ext == '.jpeg' or ext == '.ppm' or ext == '.jpg':
+        return Image.open(file_name)
+    elif ext == '.bin' or ext == '.raw':
+        return np.load(file_name)
+    elif ext == '.flo':
+        return readFlow(file_name).astype(np.float32)
+    elif ext == '.pfm':
+        flow = readPFM(file_name).astype(np.float32)
+        if len(flow.shape) == 2:
+            return flow
+        else:
+            return flow[:, :, :-1]
+    return []

AIGVDet/core/utils/utils.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from scipy import interpolate
+
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by 8 """
+    def __init__(self, dims, mode='sintel'):
+        self.ht, self.wd = dims[-2:]
+        pad_ht = (((self.ht // 8) + 1) * 8 - self.ht) % 8
+        pad_wd = (((self.wd // 8) + 1) * 8 - self.wd) % 8
+        if mode == 'sintel':
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
+        else:
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
+
+    def pad(self, *inputs):
+        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self,x):
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+
+def forward_interpolate(flow):
+    flow = flow.detach().cpu().numpy()
+    dx, dy = flow[0], flow[1]
+
+    ht, wd = dx.shape
+    x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht))
+
+    x1 = x0 + dx
+    y1 = y0 + dy
+    
+    x1 = x1.reshape(-1)
+    y1 = y1.reshape(-1)
+    dx = dx.reshape(-1)
+    dy = dy.reshape(-1)
+
+    valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht)
+    x1 = x1[valid]
+    y1 = y1[valid]
+    dx = dx[valid]
+    dy = dy[valid]
+
+    flow_x = interpolate.griddata(
+        (x1, y1), dx, (x0, y0), method='nearest', fill_value=0)
+
+    flow_y = interpolate.griddata(
+        (x1, y1), dy, (x0, y0), method='nearest', fill_value=0)
+
+    flow = np.stack([flow_x, flow_y], axis=0)
+    return torch.from_numpy(flow).float()
+
+
+def bilinear_sampler(img, coords, mode='bilinear', mask=False):
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1,1], dim=-1)
+    xgrid = 2*xgrid/(W-1) - 1
+    ygrid = 2*ygrid/(H-1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+    img = F.grid_sample(img, grid, align_corners=True)
+
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+
+    return img
+
+
+def coords_grid(batch, ht, wd, device):
+    coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+def upflow8(flow, mode='bilinear'):
+    new_size = (8 * flow.shape[2], 8 * flow.shape[3])
+    return  8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True)

AIGVDet/core/utils1/config.py

@@ -0,0 +1,156 @@
+import argparse
+import os
+import sys
+from abc import ABC
+from typing import Type
+
+
+class DefaultConfigs(ABC):
+    ####### base setting ######
+    gpus = [0]
+    seed = 3407
+    arch = "resnet50"
+    datasets = ["zhaolian_train"]
+    datasets_test = ["adm_res_abs_ddim20s"]
+    mode = "binary"
+    class_bal = False
+    batch_size = 64
+    loadSize = 256
+    cropSize = 224
+    epoch = "latest"
+    num_workers = 20
+    serial_batches = False
+    isTrain = True
+
+    # data augmentation
+    rz_interp = ["bilinear"]
+    # blur_prob = 0.0
+    blur_prob = 0.1
+    blur_sig = [0.5]
+    # jpg_prob = 0.0
+    jpg_prob = 0.1
+    jpg_method = ["cv2"]
+    jpg_qual = [75]
+    gray_prob = 0.0
+    aug_resize = True
+    aug_crop = True
+    aug_flip = True
+    aug_norm = True
+
+    ####### train setting ######
+    warmup = False
+    # warmup = True
+    warmup_epoch = 3
+    earlystop = True
+    earlystop_epoch = 5
+    optim = "adam"
+    new_optim = False
+    loss_freq = 400
+    save_latest_freq = 2000
+    save_epoch_freq = 20
+    continue_train = False
+    epoch_count = 1
+    last_epoch = -1
+    nepoch = 400
+    beta1 = 0.9
+    lr = 0.0001
+    init_type = "normal"
+    init_gain = 0.02
+    pretrained = True
+
+    # paths information
+    root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    dataset_root = os.path.join(root_dir, "data")
+    exp_root = os.path.join(root_dir, "data", "exp")
+    _exp_name = ""
+    exp_dir = ""
+    ckpt_dir = ""
+    logs_path = ""
+    ckpt_path = ""
+
+    @property
+    def exp_name(self):
+        return self._exp_name
+
+    @exp_name.setter
+    def exp_name(self, value: str):
+        self._exp_name = value
+        self.exp_dir: str = os.path.join(self.exp_root, self.exp_name)
+        self.ckpt_dir: str = os.path.join(self.exp_dir, "ckpt")
+        self.logs_path: str = os.path.join(self.exp_dir, "logs.txt")
+
+        os.makedirs(self.exp_dir, exist_ok=True)
+        os.makedirs(self.ckpt_dir, exist_ok=True)
+
+    def to_dict(self):
+        dic = {}
+        for fieldkey in dir(self):
+            fieldvalue = getattr(self, fieldkey)
+            if not fieldkey.startswith("__") and not callable(fieldvalue) and not fieldkey.startswith("_"):
+                dic[fieldkey] = fieldvalue
+        return dic
+
+
+def args_list2dict(arg_list: list):
+    assert len(arg_list) % 2 == 0, f"Override list has odd length: {arg_list}; it must be a list of pairs"
+    return dict(zip(arg_list[::2], arg_list[1::2]))
+
+
+def str2bool(v: str) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif v.lower() in ("true", "yes", "on", "y", "t", "1"):
+        return True
+    elif v.lower() in ("false", "no", "off", "n", "f", "0"):
+        return False
+    else:
+        return bool(v)
+
+
+def str2list(v: str, element_type=None) -> list:
+    if not isinstance(v, (list, tuple, set)):
+        v = v.lstrip("[").rstrip("]")
+        v = v.split(",")
+        v = list(map(str.strip, v))
+        if element_type is not None:
+            v = list(map(element_type, v))
+    return v
+
+
+CONFIGCLASS = Type[DefaultConfigs]
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpus", default=[0], type=int, nargs="+")
+parser.add_argument("--exp_name", default="", type=str)
+parser.add_argument("--ckpt", default="model_epoch_latest.pth", type=str)
+parser.add_argument("opts", default=[], nargs=argparse.REMAINDER)
+args = parser.parse_args()
+
+if os.path.exists(os.path.join(DefaultConfigs.exp_root, args.exp_name, "config.py")):
+    sys.path.insert(0, os.path.join(DefaultConfigs.exp_root, args.exp_name))
+    from config import cfg
+
+    cfg: CONFIGCLASS
+else:
+    cfg = DefaultConfigs()
+
+if args.opts:
+    opts = args_list2dict(args.opts)
+    for k, v in opts.items():
+        if not hasattr(cfg, k):
+            raise ValueError(f"Unrecognized option: {k}")
+        original_type = type(getattr(cfg, k))
+        if original_type == bool:
+            setattr(cfg, k, str2bool(v))
+        elif original_type in (list, tuple, set):
+            setattr(cfg, k, str2list(v, type(getattr(cfg, k)[0])))
+        else:
+            setattr(cfg, k, original_type(v))
+
+cfg.gpus: list = args.gpus
+os.environ["CUDA_VISIBLE_DEVICES"] = ", ".join([str(gpu) for gpu in cfg.gpus])
+cfg.exp_name = args.exp_name
+cfg.ckpt_path: str = os.path.join(cfg.ckpt_dir, args.ckpt)
+
+if isinstance(cfg.datasets, str):
+    cfg.datasets = cfg.datasets.split(",")

AIGVDet/core/utils1/datasets.py

@@ -0,0 +1,178 @@
+import os
+from io import BytesIO
+from random import choice, random
+
+import cv2
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+from PIL import Image, ImageFile
+from scipy.ndimage import gaussian_filter
+from torch.utils.data.sampler import WeightedRandomSampler
+
+from utils1.config import CONFIGCLASS
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def dataset_folder(root: str, cfg: CONFIGCLASS):
+    if cfg.mode == "binary":
+        return binary_dataset(root, cfg)
+    if cfg.mode == "filename":
+        return FileNameDataset(root, cfg)
+    raise ValueError("cfg.mode needs to be binary or filename.")
+
+
+def binary_dataset(root: str, cfg: CONFIGCLASS):
+    identity_transform = transforms.Lambda(lambda img: img)
+    
+    rz_func = identity_transform
+    
+    if cfg.isTrain:
+        crop_func = transforms.RandomCrop((448,448))
+    else:
+        crop_func = transforms.CenterCrop((448,448)) if cfg.aug_crop else identity_transform
+
+    if cfg.isTrain and cfg.aug_flip:
+        flip_func = transforms.RandomHorizontalFlip()
+    else:
+        flip_func = identity_transform
+    
+
+    return datasets.ImageFolder(
+        root,
+        transforms.Compose(
+            [
+                rz_func,
+                #change
+                transforms.Lambda(lambda img: blur_jpg_augment(img, cfg)),
+                crop_func,
+                flip_func,
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+                if cfg.aug_norm
+                else identity_transform,
+            ]
+        )
+    )
+
+
+class FileNameDataset(datasets.ImageFolder):
+    def name(self):
+        return 'FileNameDataset'
+
+    def __init__(self, opt, root):
+        self.opt = opt
+        super().__init__(root)
+
+    def __getitem__(self, index):
+        # Loading sample
+        path, target = self.samples[index]
+        return path
+
+
+def blur_jpg_augment(img: Image.Image, cfg: CONFIGCLASS):
+    img: np.ndarray = np.array(img)
+    if cfg.isTrain:
+        if random() < cfg.blur_prob:
+            sig = sample_continuous(cfg.blur_sig)
+            gaussian_blur(img, sig)
+
+        if random() < cfg.jpg_prob:
+            method = sample_discrete(cfg.jpg_method)
+            qual = sample_discrete(cfg.jpg_qual)
+            img = jpeg_from_key(img, qual, method)
+
+    return Image.fromarray(img)
+
+
+def sample_continuous(s: list):
+    if len(s) == 1:
+        return s[0]
+    if len(s) == 2:
+        rg = s[1] - s[0]
+        return random() * rg + s[0]
+    raise ValueError("Length of iterable s should be 1 or 2.")
+
+
+def sample_discrete(s: list):
+    return s[0] if len(s) == 1 else choice(s)
+
+
+def gaussian_blur(img: np.ndarray, sigma: float):
+    gaussian_filter(img[:, :, 0], output=img[:, :, 0], sigma=sigma)
+    gaussian_filter(img[:, :, 1], output=img[:, :, 1], sigma=sigma)
+    gaussian_filter(img[:, :, 2], output=img[:, :, 2], sigma=sigma)
+
+
+def cv2_jpg(img: np.ndarray, compress_val: int) -> np.ndarray:
+    img_cv2 = img[:, :, ::-1]
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), compress_val]
+    result, encimg = cv2.imencode(".jpg", img_cv2, encode_param)
+    decimg = cv2.imdecode(encimg, 1)
+    return decimg[:, :, ::-1]
+
+
+def pil_jpg(img: np.ndarray, compress_val: int):
+    out = BytesIO()
+    img = Image.fromarray(img)
+    img.save(out, format="jpeg", quality=compress_val)
+    img = Image.open(out)
+    # load from memory before ByteIO closes
+    img = np.array(img)
+    out.close()
+    return img
+
+
+jpeg_dict = {"cv2": cv2_jpg, "pil": pil_jpg}
+
+
+def jpeg_from_key(img: np.ndarray, compress_val: int, key: str) -> np.ndarray:
+    method = jpeg_dict[key]
+    return method(img, compress_val)
+
+
+rz_dict = {'bilinear': Image.BILINEAR,
+           'bicubic': Image.BICUBIC,
+           'lanczos': Image.LANCZOS,
+           'nearest': Image.NEAREST}
+def custom_resize(img: Image.Image, cfg: CONFIGCLASS) -> Image.Image:
+    interp = sample_discrete(cfg.rz_interp)
+    return TF.resize(img, cfg.loadSize, interpolation=rz_dict[interp])
+
+
+def get_dataset(cfg: CONFIGCLASS):
+    dset_lst = []
+    for dataset in cfg.datasets:
+        root = os.path.join(cfg.dataset_root, dataset)
+        dset = dataset_folder(root, cfg)
+        dset_lst.append(dset)
+    return torch.utils.data.ConcatDataset(dset_lst)
+
+
+def get_bal_sampler(dataset: torch.utils.data.ConcatDataset):
+    targets = []
+    for d in dataset.datasets:
+        targets.extend(d.targets)
+
+    ratio = np.bincount(targets)
+    w = 1.0 / torch.tensor(ratio, dtype=torch.float)
+    sample_weights = w[targets]
+    return WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights))
+
+
+def create_dataloader(cfg: CONFIGCLASS):
+    shuffle = not cfg.serial_batches if (cfg.isTrain and not cfg.class_bal) else False
+    dataset = get_dataset(cfg)
+    sampler = get_bal_sampler(dataset) if cfg.class_bal else None
+
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=cfg.batch_size,
+        shuffle=shuffle,
+        sampler=sampler,
+        num_workers=int(cfg.num_workers),
+    )

AIGVDet/core/utils1/earlystop.py

@@ -0,0 +1,46 @@
+import numpy as np
+
+from utils1.trainer import Trainer
+
+
+class EarlyStopping:
+    """Early stops the training if validation loss doesn't improve after a given patience."""
+
+    def __init__(self, patience=1, verbose=False, delta=0):
+        """
+        Args:
+            patience (int): How long to wait after last time validation loss improved.
+                            Default: 7
+            verbose (bool): If True, prints a message for each validation loss improvement.
+                            Default: False
+            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
+                            Default: 0
+        """
+        self.patience = patience
+        self.verbose = verbose
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+        self.score_max = -np.Inf
+        self.delta = delta
+
+    def __call__(self, score: float, trainer: Trainer):
+        if self.best_score is None:
+            self.best_score = score
+            self.save_checkpoint(score, trainer)
+        elif score < self.best_score - self.delta:
+            self.counter += 1
+            print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.save_checkpoint(score, trainer)
+            self.counter = 0
+
+    def save_checkpoint(self, score: float, trainer: Trainer):
+        """Saves model when validation loss decrease."""
+        if self.verbose:
+            print(f"Validation accuracy increased ({self.score_max:.6f} --> {score:.6f}).  Saving model ...")
+        trainer.save_networks("best")
+        self.score_max = score

AIGVDet/core/utils1/eval.py

@@ -0,0 +1,66 @@
+import math
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+
+from utils1.config import CONFIGCLASS
+from utils1.utils import to_cuda
+
+
+def get_val_cfg(cfg: CONFIGCLASS, split="val", copy=True):
+    if copy:
+        from copy import deepcopy
+
+        val_cfg = deepcopy(cfg)
+    else:
+        val_cfg = cfg
+    val_cfg.dataset_root = os.path.join(val_cfg.dataset_root, split)
+    val_cfg.datasets = cfg.datasets_test
+    val_cfg.isTrain = False
+    # val_cfg.aug_resize = False
+    # val_cfg.aug_crop = False
+    val_cfg.aug_flip = False
+    val_cfg.serial_batches = True
+    val_cfg.jpg_method = ["pil"]
+    # Currently assumes jpg_prob, blur_prob 0 or 1
+    if len(val_cfg.blur_sig) == 2:
+        b_sig = val_cfg.blur_sig
+        val_cfg.blur_sig = [(b_sig[0] + b_sig[1]) / 2]
+    if len(val_cfg.jpg_qual) != 1:
+        j_qual = val_cfg.jpg_qual
+        val_cfg.jpg_qual = [int((j_qual[0] + j_qual[-1]) / 2)]
+    return val_cfg
+
+def validate(model: nn.Module, cfg: CONFIGCLASS):
+    from sklearn.metrics import accuracy_score, average_precision_score, roc_auc_score
+
+    from utils1.datasets import create_dataloader
+
+    data_loader = create_dataloader(cfg)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    with torch.no_grad():
+        y_true, y_pred = [], []
+        for data in data_loader:
+            img, label, meta = data if len(data) == 3 else (*data, None)
+            in_tens = to_cuda(img, device)
+            meta = to_cuda(meta, device)
+            predict = model(in_tens, meta).sigmoid()
+            y_pred.extend(predict.flatten().tolist())
+            y_true.extend(label.flatten().tolist())
+
+    y_true, y_pred = np.array(y_true), np.array(y_pred)
+    r_acc = accuracy_score(y_true[y_true == 0], y_pred[y_true == 0] > 0.5)
+    f_acc = accuracy_score(y_true[y_true == 1], y_pred[y_true == 1] > 0.5)
+    acc = accuracy_score(y_true, y_pred > 0.5)
+    ap = average_precision_score(y_true, y_pred)
+    results = {
+        "ACC": acc,
+        "AP": ap,
+        "R_ACC": r_acc,
+        "F_ACC": f_acc,
+    }
+    return results

AIGVDet/core/utils1/trainer.py

@@ -0,0 +1,163 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from utils1.config import CONFIGCLASS
+from utils1.utils import get_network
+from utils1.warmup import GradualWarmupScheduler
+
+
+class BaseModel(nn.Module):
+    def __init__(self, cfg: CONFIGCLASS):
+        super().__init__()
+        self.cfg = cfg
+        self.total_steps = 0
+        self.isTrain = cfg.isTrain
+        self.save_dir = cfg.ckpt_dir
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.model:nn.Module
+        self.model=nn.Module.to(self.device)
+        # self.model.to(self.device)
+        #self.model.load_state_dict(torch.load('./checkpoints/optical.pth'))
+        self.optimizer: torch.optim.Optimizer
+
+    def save_networks(self, epoch: int):
+        save_filename = f"model_epoch_{epoch}.pth"
+        save_path = os.path.join(self.save_dir, save_filename)
+
+        # serialize model and optimizer to dict
+        state_dict = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "total_steps": self.total_steps,
+        }
+
+        torch.save(state_dict, save_path)
+
+    # load models from the disk
+    def load_networks(self, epoch: int):
+        load_filename = f"model_epoch_{epoch}.pth"
+        load_path = os.path.join(self.save_dir, load_filename)
+
+        if epoch==0:
+            # load_filename = f"lsun_adm.pth"
+            load_path="checkpoints/optical.pth"
+            print("loading optical path")
+        else :
+            print(f"loading the model from {load_path}")
+        
+        # print(f"loading the model from {load_path}")
+
+        # if you are using PyTorch newer than 0.4 (e.g., built from
+        # GitHub source), you can remove str() on self.device
+        state_dict = torch.load(load_path, map_location=self.device)
+        if hasattr(state_dict, "_metadata"):
+            del state_dict._metadata
+
+        self.model.load_state_dict(state_dict["model"])
+        self.total_steps = state_dict["total_steps"]
+
+        if self.isTrain and not self.cfg.new_optim:
+            self.optimizer.load_state_dict(state_dict["optimizer"])
+            # move optimizer state to GPU
+            for state in self.optimizer.state.values():
+                for k, v in state.items():
+                    if torch.is_tensor(v):
+                        state[k] = v.to(self.device)
+
+            for g in self.optimizer.param_groups:
+                g["lr"] = self.cfg.lr
+
+    def eval(self):
+        self.model.eval()
+
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+
+
+def init_weights(net: nn.Module, init_type="normal", gain=0.02):
+    def init_func(m: nn.Module):
+        classname = m.__class__.__name__
+        if hasattr(m, "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):
+            if init_type == "normal":
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == "xavier":
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == "kaiming":
+                init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+            elif init_type == "orthogonal":
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError(f"initialization method [{init_type}] is not implemented")
+            if hasattr(m, "bias") and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find("BatchNorm2d") != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print(f"initialize network with {init_type}")
+    net.apply(init_func)
+
+
+class Trainer(BaseModel):
+    def name(self):
+        return "Trainer"
+
+    def __init__(self, cfg: CONFIGCLASS):
+        super().__init__(cfg)
+        self.arch = cfg.arch
+        self.model = get_network(self.arch, cfg.isTrain, cfg.continue_train, cfg.init_gain, cfg.pretrained)
+
+        self.loss_fn = nn.BCEWithLogitsLoss()
+        # initialize optimizers
+        if cfg.optim == "adam":
+            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=cfg.lr, betas=(cfg.beta1, 0.999))
+        elif cfg.optim == "sgd":
+            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=cfg.lr, momentum=0.9, weight_decay=5e-4)
+        else:
+            raise ValueError("optim should be [adam, sgd]")
+        if cfg.warmup:
+            scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+                self.optimizer, cfg.nepoch - cfg.warmup_epoch, eta_min=1e-6
+            )
+            self.scheduler = GradualWarmupScheduler(
+                self.optimizer, multiplier=1, total_epoch=cfg.warmup_epoch, after_scheduler=scheduler_cosine
+            )
+            self.scheduler.step()
+        if cfg.continue_train:
+            self.load_networks(cfg.epoch)
+        self.model.to(self.device)
+
+        
+
+    def adjust_learning_rate(self, min_lr=1e-6):
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] /= 10.0
+            if param_group["lr"] < min_lr:
+                return False
+        return True
+
+    def set_input(self, input):
+        img, label, meta = input if len(input) == 3 else (input[0], input[1], {})
+        self.input = img.to(self.device)
+        self.label = label.to(self.device).float()
+        for k in meta.keys():
+            if isinstance(meta[k], torch.Tensor):
+                meta[k] = meta[k].to(self.device)
+        self.meta = meta
+
+    def forward(self):
+        self.output = self.model(self.input, self.meta)
+
+    def get_loss(self):
+        return self.loss_fn(self.output.squeeze(1), self.label)
+
+    def optimize_parameters(self):
+        self.forward()
+        self.loss = self.loss_fn(self.output.squeeze(1), self.label)
+        self.optimizer.zero_grad()
+        self.loss.backward()
+        self.optimizer.step()

AIGVDet/core/utils1/utils.py

@@ -0,0 +1,109 @@
+import argparse
+import os
+import sys
+import time
+import warnings
+from importlib import import_module
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.nn.functional")
+
+
+def str2bool(v: str, strict=True) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif isinstance(v, str):
+        if v.lower() in ("true", "yes", "on" "t", "y", "1"):
+            return True
+        elif v.lower() in ("false", "no", "off", "f", "n", "0"):
+            return False
+    if strict:
+        raise argparse.ArgumentTypeError("Unsupported value encountered.")
+    else:
+        return True
+
+
+def to_cuda(data, device="cuda", exclude_keys: "list[str]" = None):
+    if isinstance(data, torch.Tensor):
+        data = data.to(device)
+    elif isinstance(data, (tuple, list, set)):
+        data = [to_cuda(b, device) for b in data]
+    elif isinstance(data, dict):
+        if exclude_keys is None:
+            exclude_keys = []
+        for k in data.keys():
+            if k not in exclude_keys:
+                data[k] = to_cuda(data[k], device)
+    else:
+        # raise TypeError(f"Unsupported type: {type(data)}")
+        data = data
+    return data
+
+
+class HiddenPrints:
+    def __enter__(self):
+        self._original_stdout = sys.stdout
+        sys.stdout = open(os.devnull, "w")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout.close()
+        sys.stdout = self._original_stdout
+
+
+class Logger(object):
+    def __init__(self):
+        self.terminal = sys.stdout
+        self.file = None
+
+    def open(self, file, mode=None):
+        if mode is None:
+            mode = "w"
+        self.file = open(file, mode)
+
+    def write(self, message, is_terminal=1, is_file=1):
+        if "\r" in message:
+            is_file = 0
+        if is_terminal == 1:
+            self.terminal.write(message)
+            self.terminal.flush()
+        if is_file == 1:
+            self.file.write(message)
+            self.file.flush()
+
+    def flush(self):
+        # this flush method is needed for python 3 compatibility.
+        # this handles the flush command by doing nothing.
+        # you might want to specify some extra behavior here.
+        pass
+
+
+def get_network(arch: str, isTrain=False, continue_train=False, init_gain=0.02, pretrained=True):
+    if "resnet" in arch:
+        from ...networks.resnet import ResNet
+
+        resnet = getattr(import_module("...networks.resnet", package=__package__), arch)
+        if isTrain:
+            if continue_train:
+                model: ResNet = resnet(num_classes=1)
+            else:
+                model: ResNet = resnet(pretrained=pretrained)
+                model.fc = nn.Linear(2048, 1)
+                nn.init.normal_(model.fc.weight.data, 0.0, init_gain)
+        else:
+            model: ResNet = resnet(num_classes=1)
+        return model
+    else:
+        raise ValueError(f"Unsupported arch: {arch}")
+
+
+def pad_img_to_square(img: np.ndarray):
+    H, W = img.shape[:2]
+    if H != W:
+        new_size = max(H, W)
+        img = np.pad(img, ((0, new_size - H), (0, new_size - W), (0, 0)), mode="constant")
+        assert img.shape[0] == img.shape[1] == new_size
+    return img

AIGVDet/core/utils1/utils1/config.py

@@ -0,0 +1,157 @@
+import argparse
+import os
+import sys
+from abc import ABC
+from typing import Type
+
+
+class DefaultConfigs(ABC):
+    ####### base setting ######
+    gpus = [0]
+    seed = 3407
+    arch = "resnet50"
+    datasets = ["zhaolian_train"]
+    datasets_test = ["adm_res_abs_ddim20s"]
+    mode = "binary"
+    class_bal = False
+    batch_size = 64
+    loadSize = 256
+    cropSize = 224
+    epoch = "latest"
+    num_workers = 20
+    serial_batches = False
+    isTrain = True
+
+    # data augmentation
+    rz_interp = ["bilinear"]
+    # blur_prob = 0.0
+    blur_prob = 0.1
+    blur_sig = [0.5]
+    # jpg_prob = 0.0
+    jpg_prob = 0.1
+    jpg_method = ["cv2"]
+    jpg_qual = [75]
+    gray_prob = 0.0
+    aug_resize = True
+    aug_crop = True
+    aug_flip = True
+    aug_norm = True
+
+    ####### train setting ######
+    warmup = False
+    # warmup = True
+    warmup_epoch = 3
+    earlystop = True
+    earlystop_epoch = 5
+    optim = "adam"
+    new_optim = False
+    loss_freq = 400
+    save_latest_freq = 2000
+    save_epoch_freq = 20
+    continue_train = False
+    epoch_count = 1
+    last_epoch = -1
+    nepoch = 400
+    beta1 = 0.9
+    lr = 0.0001
+    init_type = "normal"
+    init_gain = 0.02
+    pretrained = True
+
+    # paths information
+    root_dir1 = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    root_dir = os.path.dirname(root_dir1)
+    dataset_root = os.path.join(root_dir, "data")
+    exp_root = os.path.join(root_dir, "data", "exp")
+    _exp_name = ""
+    exp_dir = ""
+    ckpt_dir = ""
+    logs_path = ""
+    ckpt_path = ""
+
+    @property
+    def exp_name(self):
+        return self._exp_name
+
+    @exp_name.setter
+    def exp_name(self, value: str):
+        self._exp_name = value
+        self.exp_dir: str = os.path.join(self.exp_root, self.exp_name)
+        self.ckpt_dir: str = os.path.join(self.exp_dir, "ckpt")
+        self.logs_path: str = os.path.join(self.exp_dir, "logs.txt")
+
+        os.makedirs(self.exp_dir, exist_ok=True)
+        os.makedirs(self.ckpt_dir, exist_ok=True)
+
+    def to_dict(self):
+        dic = {}
+        for fieldkey in dir(self):
+            fieldvalue = getattr(self, fieldkey)
+            if not fieldkey.startswith("__") and not callable(fieldvalue) and not fieldkey.startswith("_"):
+                dic[fieldkey] = fieldvalue
+        return dic
+
+
+def args_list2dict(arg_list: list):
+    assert len(arg_list) % 2 == 0, f"Override list has odd length: {arg_list}; it must be a list of pairs"
+    return dict(zip(arg_list[::2], arg_list[1::2]))
+
+
+def str2bool(v: str) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif v.lower() in ("true", "yes", "on", "y", "t", "1"):
+        return True
+    elif v.lower() in ("false", "no", "off", "n", "f", "0"):
+        return False
+    else:
+        return bool(v)
+
+
+def str2list(v: str, element_type=None) -> list:
+    if not isinstance(v, (list, tuple, set)):
+        v = v.lstrip("[").rstrip("]")
+        v = v.split(",")
+        v = list(map(str.strip, v))
+        if element_type is not None:
+            v = list(map(element_type, v))
+    return v
+
+
+CONFIGCLASS = Type[DefaultConfigs]
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--gpus", default=[0], type=int, nargs="+")
+parser.add_argument("--exp_name", default="", type=str)
+parser.add_argument("--ckpt", default="model_epoch_latest.pth", type=str)
+parser.add_argument("opts", default=[], nargs=argparse.REMAINDER)
+args = parser.parse_args()
+
+if os.path.exists(os.path.join(DefaultConfigs.exp_root, args.exp_name, "config.py")):
+    sys.path.insert(0, os.path.join(DefaultConfigs.exp_root, args.exp_name))
+    from config import cfg
+
+    cfg: CONFIGCLASS
+else:
+    cfg = DefaultConfigs()
+
+if args.opts:
+    opts = args_list2dict(args.opts)
+    for k, v in opts.items():
+        if not hasattr(cfg, k):
+            raise ValueError(f"Unrecognized option: {k}")
+        original_type = type(getattr(cfg, k))
+        if original_type == bool:
+            setattr(cfg, k, str2bool(v))
+        elif original_type in (list, tuple, set):
+            setattr(cfg, k, str2list(v, type(getattr(cfg, k)[0])))
+        else:
+            setattr(cfg, k, original_type(v))
+
+cfg.gpus: list = args.gpus
+os.environ["CUDA_VISIBLE_DEVICES"] = ", ".join([str(gpu) for gpu in cfg.gpus])
+cfg.exp_name = args.exp_name
+cfg.ckpt_path: str = os.path.join(cfg.ckpt_dir, args.ckpt)
+
+if isinstance(cfg.datasets, str):
+    cfg.datasets = cfg.datasets.split(",")

AIGVDet/core/utils1/utils1/datasets.py

@@ -0,0 +1,178 @@
+import os
+from io import BytesIO
+from random import choice, random
+
+import cv2
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+from PIL import Image, ImageFile
+from scipy.ndimage import gaussian_filter
+from torch.utils.data.sampler import WeightedRandomSampler
+
+from .config import CONFIGCLASS
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+def dataset_folder(root: str, cfg: CONFIGCLASS):
+    if cfg.mode == "binary":
+        return binary_dataset(root, cfg)
+    if cfg.mode == "filename":
+        return FileNameDataset(root, cfg)
+    raise ValueError("cfg.mode needs to be binary or filename.")
+
+
+def binary_dataset(root: str, cfg: CONFIGCLASS):
+    identity_transform = transforms.Lambda(lambda img: img)
+    
+    rz_func = identity_transform
+    
+    if cfg.isTrain:
+        crop_func = transforms.RandomCrop((448,448))
+    else:
+        crop_func = transforms.CenterCrop((448,448)) if cfg.aug_crop else identity_transform
+
+    if cfg.isTrain and cfg.aug_flip:
+        flip_func = transforms.RandomHorizontalFlip()
+    else:
+        flip_func = identity_transform
+    
+
+    return datasets.ImageFolder(
+        root,
+        transforms.Compose(
+            [
+                rz_func,
+                #change
+                transforms.Lambda(lambda img: blur_jpg_augment(img, cfg)),
+                crop_func,
+                flip_func,
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+                if cfg.aug_norm
+                else identity_transform,
+            ]
+        )
+    )
+
+
+class FileNameDataset(datasets.ImageFolder):
+    def name(self):
+        return 'FileNameDataset'
+
+    def __init__(self, opt, root):
+        self.opt = opt
+        super().__init__(root)
+
+    def __getitem__(self, index):
+        # Loading sample
+        path, target = self.samples[index]
+        return path
+
+
+def blur_jpg_augment(img: Image.Image, cfg: CONFIGCLASS):
+    img: np.ndarray = np.array(img)
+    if cfg.isTrain:
+        if random() < cfg.blur_prob:
+            sig = sample_continuous(cfg.blur_sig)
+            gaussian_blur(img, sig)
+
+        if random() < cfg.jpg_prob:
+            method = sample_discrete(cfg.jpg_method)
+            qual = sample_discrete(cfg.jpg_qual)
+            img = jpeg_from_key(img, qual, method)
+
+    return Image.fromarray(img)
+
+
+def sample_continuous(s: list):
+    if len(s) == 1:
+        return s[0]
+    if len(s) == 2:
+        rg = s[1] - s[0]
+        return random() * rg + s[0]
+    raise ValueError("Length of iterable s should be 1 or 2.")
+
+
+def sample_discrete(s: list):
+    return s[0] if len(s) == 1 else choice(s)
+
+
+def gaussian_blur(img: np.ndarray, sigma: float):
+    gaussian_filter(img[:, :, 0], output=img[:, :, 0], sigma=sigma)
+    gaussian_filter(img[:, :, 1], output=img[:, :, 1], sigma=sigma)
+    gaussian_filter(img[:, :, 2], output=img[:, :, 2], sigma=sigma)
+
+
+def cv2_jpg(img: np.ndarray, compress_val: int) -> np.ndarray:
+    img_cv2 = img[:, :, ::-1]
+    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), compress_val]
+    result, encimg = cv2.imencode(".jpg", img_cv2, encode_param)
+    decimg = cv2.imdecode(encimg, 1)
+    return decimg[:, :, ::-1]
+
+
+def pil_jpg(img: np.ndarray, compress_val: int):
+    out = BytesIO()
+    img = Image.fromarray(img)
+    img.save(out, format="jpeg", quality=compress_val)
+    img = Image.open(out)
+    # load from memory before ByteIO closes
+    img = np.array(img)
+    out.close()
+    return img
+
+
+jpeg_dict = {"cv2": cv2_jpg, "pil": pil_jpg}
+
+
+def jpeg_from_key(img: np.ndarray, compress_val: int, key: str) -> np.ndarray:
+    method = jpeg_dict[key]
+    return method(img, compress_val)
+
+
+rz_dict = {'bilinear': Image.BILINEAR,
+           'bicubic': Image.BICUBIC,
+           'lanczos': Image.LANCZOS,
+           'nearest': Image.NEAREST}
+def custom_resize(img: Image.Image, cfg: CONFIGCLASS) -> Image.Image:
+    interp = sample_discrete(cfg.rz_interp)
+    return TF.resize(img, cfg.loadSize, interpolation=rz_dict[interp])
+
+
+def get_dataset(cfg: CONFIGCLASS):
+    dset_lst = []
+    for dataset in cfg.datasets:
+        root = os.path.join(cfg.dataset_root, dataset)
+        dset = dataset_folder(root, cfg)
+        dset_lst.append(dset)
+    return torch.utils.data.ConcatDataset(dset_lst)
+
+
+def get_bal_sampler(dataset: torch.utils.data.ConcatDataset):
+    targets = []
+    for d in dataset.datasets:
+        targets.extend(d.targets)
+
+    ratio = np.bincount(targets)
+    w = 1.0 / torch.tensor(ratio, dtype=torch.float)
+    sample_weights = w[targets]
+    return WeightedRandomSampler(weights=sample_weights, num_samples=len(sample_weights))
+
+
+def create_dataloader(cfg: CONFIGCLASS):
+    shuffle = not cfg.serial_batches if (cfg.isTrain and not cfg.class_bal) else False
+    dataset = get_dataset(cfg)
+    sampler = get_bal_sampler(dataset) if cfg.class_bal else None
+
+    return torch.utils.data.DataLoader(
+        dataset,
+        batch_size=cfg.batch_size,
+        shuffle=shuffle,
+        sampler=sampler,
+        num_workers=int(cfg.num_workers),
+    )

AIGVDet/core/utils1/utils1/earlystop.py

@@ -0,0 +1,46 @@
+import numpy as np
+
+from .trainer import Trainer
+
+
+class EarlyStopping:
+    """Early stops the training if validation loss doesn't improve after a given patience."""
+
+    def __init__(self, patience=1, verbose=False, delta=0):
+        """
+        Args:
+            patience (int): How long to wait after last time validation loss improved.
+                            Default: 7
+            verbose (bool): If True, prints a message for each validation loss improvement.
+                            Default: False
+            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
+                            Default: 0
+        """
+        self.patience = patience
+        self.verbose = verbose
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+        self.score_max = -np.Inf
+        self.delta = delta
+
+    def __call__(self, score: float, trainer: Trainer):
+        if self.best_score is None:
+            self.best_score = score
+            self.save_checkpoint(score, trainer)
+        elif score < self.best_score - self.delta:
+            self.counter += 1
+            print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.save_checkpoint(score, trainer)
+            self.counter = 0
+
+    def save_checkpoint(self, score: float, trainer: Trainer):
+        """Saves model when validation loss decrease."""
+        if self.verbose:
+            print(f"Validation accuracy increased ({self.score_max:.6f} --> {score:.6f}).  Saving model ...")
+        trainer.save_networks("best")
+        self.score_max = score

AIGVDet/core/utils1/utils1/eval.py

@@ -0,0 +1,66 @@
+import math
+import os
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+
+from .config import CONFIGCLASS
+from .utils import to_cuda
+
+
+def get_val_cfg(cfg: CONFIGCLASS, split="val", copy=True):
+    if copy:
+        from copy import deepcopy
+
+        val_cfg = deepcopy(cfg)
+    else:
+        val_cfg = cfg
+    val_cfg.dataset_root = os.path.join(val_cfg.dataset_root, split)
+    val_cfg.datasets = cfg.datasets_test
+    val_cfg.isTrain = False
+    # val_cfg.aug_resize = False
+    # val_cfg.aug_crop = False
+    val_cfg.aug_flip = False
+    val_cfg.serial_batches = True
+    val_cfg.jpg_method = ["pil"]
+    # Currently assumes jpg_prob, blur_prob 0 or 1
+    if len(val_cfg.blur_sig) == 2:
+        b_sig = val_cfg.blur_sig
+        val_cfg.blur_sig = [(b_sig[0] + b_sig[1]) / 2]
+    if len(val_cfg.jpg_qual) != 1:
+        j_qual = val_cfg.jpg_qual
+        val_cfg.jpg_qual = [int((j_qual[0] + j_qual[-1]) / 2)]
+    return val_cfg
+
+def validate(model: nn.Module, cfg: CONFIGCLASS):
+    from sklearn.metrics import accuracy_score, average_precision_score, roc_auc_score
+
+    from .datasets import create_dataloader
+
+    data_loader = create_dataloader(cfg)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    with torch.no_grad():
+        y_true, y_pred = [], []
+        for data in data_loader:
+            img, label, meta = data if len(data) == 3 else (*data, None)
+            in_tens = to_cuda(img, device)
+            meta = to_cuda(meta, device)
+            predict = model(in_tens, meta).sigmoid()
+            y_pred.extend(predict.flatten().tolist())
+            y_true.extend(label.flatten().tolist())
+
+    y_true, y_pred = np.array(y_true), np.array(y_pred)
+    r_acc = accuracy_score(y_true[y_true == 0], y_pred[y_true == 0] > 0.5)
+    f_acc = accuracy_score(y_true[y_true == 1], y_pred[y_true == 1] > 0.5)
+    acc = accuracy_score(y_true, y_pred > 0.5)
+    ap = average_precision_score(y_true, y_pred)
+    results = {
+        "ACC": acc,
+        "AP": ap,
+        "R_ACC": r_acc,
+        "F_ACC": f_acc,
+    }
+    return results

AIGVDet/core/utils1/utils1/trainer.py

@@ -0,0 +1,169 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.nn import init
+
+from .config import CONFIGCLASS
+from .utils import get_network
+from .warmup import GradualWarmupScheduler
+
+
+class BaseModel(nn.Module):
+    def __init__(self, cfg: CONFIGCLASS):
+        super().__init__()
+        self.cfg = cfg
+        self.total_steps = 0
+        self.isTrain = cfg.isTrain
+        self.save_dir = cfg.ckpt_dir
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.model:nn.Module
+        self.model=nn.Module.to(self.device)
+        # self.model.to(self.device)
+        self.model.load_state_dict(torch.load('./checkpoints/optical.pth'))
+        self.optimizer: torch.optim.Optimizer
+
+    def save_networks(self, epoch: int):
+        save_filename = f"model_epoch_{epoch}.pth"
+        save_path = os.path.join(self.save_dir, save_filename)
+
+        # serialize model and optimizer to dict
+        state_dict = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "total_steps": self.total_steps,
+        }
+
+        torch.save(state_dict, save_path)
+
+    # load models from the disk
+    def load_networks(self, epoch: int):
+        load_filename = f"model_epoch_{epoch}.pth"
+        load_path = os.path.join(self.save_dir, load_filename)
+
+        if epoch==0:
+            # load_filename = f"lsun_adm.pth"
+            load_path="checkpoints/optical.pth"
+            print("loading optical path")
+        else :
+            print(f"loading the model from {load_path}")
+        
+        # print(f"loading the model from {load_path}")
+
+        # if you are using PyTorch newer than 0.4 (e.g., built from
+        # GitHub source), you can remove str() on self.device
+        state_dict = torch.load(load_path, map_location=self.device)
+        if hasattr(state_dict, "_metadata"):
+            del state_dict._metadata
+
+        self.model.load_state_dict(state_dict["model"])
+        self.total_steps = state_dict["total_steps"]
+
+        if self.isTrain and not self.cfg.new_optim:
+            self.optimizer.load_state_dict(state_dict["optimizer"])
+            # move optimizer state to GPU
+            for state in self.optimizer.state.values():
+                for k, v in state.items():
+                    if torch.is_tensor(v):
+                        state[k] = v.to(self.device)
+
+            for g in self.optimizer.param_groups:
+                g["lr"] = self.cfg.lr
+
+    def eval(self):
+        self.model.eval()
+
+    def test(self):
+        with torch.no_grad():
+            self.forward()
+
+
+def init_weights(net: nn.Module, init_type="normal", gain=0.02):
+    def init_func(m: nn.Module):
+        classname = m.__class__.__name__
+        if hasattr(m, "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):
+            if init_type == "normal":
+                init.normal_(m.weight.data, 0.0, gain)
+            elif init_type == "xavier":
+                init.xavier_normal_(m.weight.data, gain=gain)
+            elif init_type == "kaiming":
+                init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
+            elif init_type == "orthogonal":
+                init.orthogonal_(m.weight.data, gain=gain)
+            else:
+                raise NotImplementedError(f"initialization method [{init_type}] is not implemented")
+            if hasattr(m, "bias") and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find("BatchNorm2d") != -1:
+            init.normal_(m.weight.data, 1.0, gain)
+            init.constant_(m.bias.data, 0.0)
+
+    print(f"initialize network with {init_type}")
+    net.apply(init_func)
+
+
+class Trainer(BaseModel):
+    def name(self):
+        return "Trainer"
+
+    def __init__(self, cfg: CONFIGCLASS):
+        super().__init__(cfg)
+        self.arch = cfg.arch
+        self.model = get_network(self.arch, cfg.isTrain, cfg.continue_train, cfg.init_gain, cfg.pretrained)
+
+        self.loss_fn = nn.BCEWithLogitsLoss()
+        # initialize optimizers
+        if cfg.optim == "adam":
+            self.optimizer = torch.optim.Adam(self.model.parameters(), lr=cfg.lr, betas=(cfg.beta1, 0.999))
+        elif cfg.optim == "sgd":
+            self.optimizer = torch.optim.SGD(self.model.parameters(), lr=cfg.lr, momentum=0.9, weight_decay=5e-4)
+        else:
+            raise ValueError("optim should be [adam, sgd]")
+        if cfg.warmup:
+            scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR(
+                self.optimizer, cfg.nepoch - cfg.warmup_epoch, eta_min=1e-6
+            )
+            self.scheduler = GradualWarmupScheduler(
+                self.optimizer, multiplier=1, total_epoch=cfg.warmup_epoch, after_scheduler=scheduler_cosine
+            )
+            self.scheduler.step()
+        if cfg.continue_train:
+            self.load_networks(cfg.epoch)
+        self.model.to(self.device)
+
+        # self.model.load_state_dict(torch.load('checkpoints/optical.pth'))
+        load_path='checkpoints/optical.pth'
+        state_dict = torch.load(load_path, map_location=self.device)
+
+
+        self.model.load_state_dict(state_dict["model"])
+        
+
+    def adjust_learning_rate(self, min_lr=1e-6):
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] /= 10.0
+            if param_group["lr"] < min_lr:
+                return False
+        return True
+
+    def set_input(self, input):
+        img, label, meta = input if len(input) == 3 else (input[0], input[1], {})
+        self.input = img.to(self.device)
+        self.label = label.to(self.device).float()
+        for k in meta.keys():
+            if isinstance(meta[k], torch.Tensor):
+                meta[k] = meta[k].to(self.device)
+        self.meta = meta
+
+    def forward(self):
+        self.output = self.model(self.input, self.meta)
+
+    def get_loss(self):
+        return self.loss_fn(self.output.squeeze(1), self.label)
+
+    def optimize_parameters(self):
+        self.forward()
+        self.loss = self.loss_fn(self.output.squeeze(1), self.label)
+        self.optimizer.zero_grad()
+        self.loss.backward()
+        self.optimizer.step()

AIGVDet/core/utils1/utils1/utils.py

@@ -0,0 +1,109 @@
+import argparse
+import os
+import sys
+import time
+import warnings
+from importlib import import_module
+
+import numpy as np
+import torch
+import torch.nn as nn
+from PIL import Image
+
+warnings.filterwarnings("ignore", category=UserWarning, module="torch.nn.functional")
+
+
+def str2bool(v: str, strict=True) -> bool:
+    if isinstance(v, bool):
+        return v
+    elif isinstance(v, str):
+        if v.lower() in ("true", "yes", "on" "t", "y", "1"):
+            return True
+        elif v.lower() in ("false", "no", "off", "f", "n", "0"):
+            return False
+    if strict:
+        raise argparse.ArgumentTypeError("Unsupported value encountered.")
+    else:
+        return True
+
+
+def to_cuda(data, device="cuda", exclude_keys: "list[str]" = None):
+    if isinstance(data, torch.Tensor):
+        data = data.to(device)
+    elif isinstance(data, (tuple, list, set)):
+        data = [to_cuda(b, device) for b in data]
+    elif isinstance(data, dict):
+        if exclude_keys is None:
+            exclude_keys = []
+        for k in data.keys():
+            if k not in exclude_keys:
+                data[k] = to_cuda(data[k], device)
+    else:
+        # raise TypeError(f"Unsupported type: {type(data)}")
+        data = data
+    return data
+
+
+class HiddenPrints:
+    def __enter__(self):
+        self._original_stdout = sys.stdout
+        sys.stdout = open(os.devnull, "w")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout.close()
+        sys.stdout = self._original_stdout
+
+
+class Logger(object):
+    def __init__(self):
+        self.terminal = sys.stdout
+        self.file = None
+
+    def open(self, file, mode=None):
+        if mode is None:
+            mode = "w"
+        self.file = open(file, mode)
+
+    def write(self, message, is_terminal=1, is_file=1):
+        if "\r" in message:
+            is_file = 0
+        if is_terminal == 1:
+            self.terminal.write(message)
+            self.terminal.flush()
+        if is_file == 1:
+            self.file.write(message)
+            self.file.flush()
+
+    def flush(self):
+        # this flush method is needed for python 3 compatibility.
+        # this handles the flush command by doing nothing.
+        # you might want to specify some extra behavior here.
+        pass
+
+
+def get_network(arch: str, isTrain=False, continue_train=False, init_gain=0.02, pretrained=True):
+    if "resnet" in arch:
+        from networks.resnet import ResNet
+
+        resnet = getattr(import_module("networks.resnet"), arch)
+        if isTrain:
+            if continue_train:
+                model: ResNet = resnet(num_classes=1)
+            else:
+                model: ResNet = resnet(pretrained=pretrained)
+                model.fc = nn.Linear(2048, 1)
+                nn.init.normal_(model.fc.weight.data, 0.0, init_gain)
+        else:
+            model: ResNet = resnet(num_classes=1)
+        return model
+    else:
+        raise ValueError(f"Unsupported arch: {arch}")
+
+
+def pad_img_to_square(img: np.ndarray):
+    H, W = img.shape[:2]
+    if H != W:
+        new_size = max(H, W)
+        img = np.pad(img, ((0, new_size - H), (0, new_size - W), (0, 0)), mode="constant")
+        assert img.shape[0] == img.shape[1] == new_size
+    return img

AIGVDet/core/utils1/utils1/warmup.py

@@ -0,0 +1,70 @@
+from torch.optim.lr_scheduler import ReduceLROnPlateau, _LRScheduler
+
+
+class GradualWarmupScheduler(_LRScheduler):
+    """Gradually warm-up(increasing) learning rate in optimizer.
+    Proposed in 'Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour'.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        multiplier: target learning rate = base lr * multiplier if multiplier > 1.0. if multiplier = 1.0, lr starts from 0 and ends up with the base_lr.
+        total_epoch: target learning rate is reached at total_epoch, gradually
+        after_scheduler: after target_epoch, use this scheduler(eg. ReduceLROnPlateau)
+    """
+
+    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):
+        self.multiplier = multiplier
+        if self.multiplier < 1.0:
+            raise ValueError("multiplier should be greater thant or equal to 1.")
+        self.total_epoch = total_epoch
+        self.after_scheduler = after_scheduler
+        self.finished = False
+        super().__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch > self.total_epoch:
+            if self.after_scheduler:
+                if not self.finished:
+                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]
+                    self.finished = True
+                return self.after_scheduler.get_last_lr()
+            return [base_lr * self.multiplier for base_lr in self.base_lrs]
+
+        if self.multiplier == 1.0:
+            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]
+        else:
+            return [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+
+    def step_ReduceLROnPlateau(self, metrics, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+        self.last_epoch = (
+            epoch if epoch != 0 else 1
+        )  # ReduceLROnPlateau is called at the end of epoch, whereas others are called at beginning
+        if self.last_epoch <= self.total_epoch:
+            warmup_lr = [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+            for param_group, lr in zip(self.optimizer.param_groups, warmup_lr):
+                param_group["lr"] = lr
+        else:
+            if epoch is None:
+                self.after_scheduler.step(metrics, None)
+            else:
+                self.after_scheduler.step(metrics, epoch - self.total_epoch)
+
+    def step(self, epoch=None, metrics=None):
+        if type(self.after_scheduler) != ReduceLROnPlateau:
+            if self.finished and self.after_scheduler:
+                if epoch is None:
+                    self.after_scheduler.step(None)
+                else:
+                    self.after_scheduler.step(epoch - self.total_epoch)
+            else:
+                return super().step(epoch)
+        else:
+            self.step_ReduceLROnPlateau(metrics, epoch)

AIGVDet/core/utils1/warmup.py

@@ -0,0 +1,70 @@
+from torch.optim.lr_scheduler import ReduceLROnPlateau, _LRScheduler
+
+
+class GradualWarmupScheduler(_LRScheduler):
+    """Gradually warm-up(increasing) learning rate in optimizer.
+    Proposed in 'Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour'.
+
+    Args:
+        optimizer (Optimizer): Wrapped optimizer.
+        multiplier: target learning rate = base lr * multiplier if multiplier > 1.0. if multiplier = 1.0, lr starts from 0 and ends up with the base_lr.
+        total_epoch: target learning rate is reached at total_epoch, gradually
+        after_scheduler: after target_epoch, use this scheduler(eg. ReduceLROnPlateau)
+    """
+
+    def __init__(self, optimizer, multiplier, total_epoch, after_scheduler=None):
+        self.multiplier = multiplier
+        if self.multiplier < 1.0:
+            raise ValueError("multiplier should be greater thant or equal to 1.")
+        self.total_epoch = total_epoch
+        self.after_scheduler = after_scheduler
+        self.finished = False
+        super().__init__(optimizer)
+
+    def get_lr(self):
+        if self.last_epoch > self.total_epoch:
+            if self.after_scheduler:
+                if not self.finished:
+                    self.after_scheduler.base_lrs = [base_lr * self.multiplier for base_lr in self.base_lrs]
+                    self.finished = True
+                return self.after_scheduler.get_last_lr()
+            return [base_lr * self.multiplier for base_lr in self.base_lrs]
+
+        if self.multiplier == 1.0:
+            return [base_lr * (float(self.last_epoch) / self.total_epoch) for base_lr in self.base_lrs]
+        else:
+            return [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+
+    def step_ReduceLROnPlateau(self, metrics, epoch=None):
+        if epoch is None:
+            epoch = self.last_epoch + 1
+        self.last_epoch = (
+            epoch if epoch != 0 else 1
+        )  # ReduceLROnPlateau is called at the end of epoch, whereas others are called at beginning
+        if self.last_epoch <= self.total_epoch:
+            warmup_lr = [
+                base_lr * ((self.multiplier - 1.0) * self.last_epoch / self.total_epoch + 1.0)
+                for base_lr in self.base_lrs
+            ]
+            for param_group, lr in zip(self.optimizer.param_groups, warmup_lr):
+                param_group["lr"] = lr
+        else:
+            if epoch is None:
+                self.after_scheduler.step(metrics, None)
+            else:
+                self.after_scheduler.step(metrics, epoch - self.total_epoch)
+
+    def step(self, epoch=None, metrics=None):
+        if type(self.after_scheduler) != ReduceLROnPlateau:
+            if self.finished and self.after_scheduler:
+                if epoch is None:
+                    self.after_scheduler.step(None)
+                else:
+                    self.after_scheduler.step(epoch - self.total_epoch)
+            else:
+                return super().step(epoch)
+        else:
+            self.step_ReduceLROnPlateau(metrics, epoch)

AIGVDet/docker-compose.yml

@@ -0,0 +1,17 @@
+version: "3.9"
+
+services:
+  web:
+    build: .
+    ports:
+      - "8003:8003"
+    restart: unless-stopped
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    ipc: host
+    shm_size: "24gb"

AIGVDet/main.py

@@ -0,0 +1,78 @@
+import os
+import json
+import shutil
+from typing import Dict, Optional
+
+from .run import RUN
+
+
+def run_video_to_json(
+    video_path: str,
+    output_json_path: Optional[str] = None,
+    model_optical_path: str = "checkpoints/optical.pth",
+    model_original_path: str = "checkpoints/original.pth",
+    frame_root: str = "frame",
+    optical_root: str = "optical_result"
+) -> Dict:
+    """
+    Xử lý 1 video và ghi kết quả ra file JSON.
+
+    Returns:
+        dict kết quả (đồng thời ghi ra JSON)
+    """
+
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    if not os.path.isabs(model_optical_path):
+        model_optical_path = os.path.join(script_dir, model_optical_path)
+    if not os.path.isabs(model_original_path):
+        model_original_path = os.path.join(script_dir, model_original_path)
+
+    results = {}
+
+    if not os.path.isfile(video_path):
+        raise FileNotFoundError(f"File not found: {video_path}")
+
+    video_name = os.path.basename(video_path)
+    video_id = os.path.splitext(video_name)[0]
+
+    folder_original = os.path.join(frame_root, video_id)
+    folder_optical = os.path.join(optical_root, video_id)
+
+    args = [
+        '--path', video_path,
+        '--folder_original_path', folder_original,
+        '--folder_optical_flow_path', folder_optical,
+        '--model_optical_flow_path', model_optical_path,
+        '--model_original_path', model_original_path
+    ]
+
+    try:
+        output = RUN(args)
+        results[video_id] = {
+            "video_name": video_name,
+            "authentic_confidence_score": round(output["real_score"], 4),
+            "synthetic_confidence_score": round(output["fake_score"], 4)
+        }
+
+    except Exception as e:
+        results[video_id] = {
+            "video_name": video_name,
+            "error": str(e)
+        }
+
+    finally:
+        # Clean up intermediate folders
+        for folder in [folder_original, folder_optical]:
+            try:
+                if os.path.exists(folder):
+                    shutil.rmtree(folder)
+            except Exception:
+                pass
+
+    # Save result to JSON or return dict
+    if output_json_path:
+        os.makedirs(os.path.dirname(output_json_path), exist_ok=True)
+        with open(output_json_path, "w", encoding="utf-8") as f:
+            json.dump(results, f, ensure_ascii=False, indent=2)
+
+    return results

AIGVDet/networks/resnet.py

@@ -0,0 +1,211 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+
+__all__ = ["ResNet", "resnet18", "resnet34", "resnet50", "resnet101", "resnet152"]
+
+
+model_urls = {
+    "resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
+    "resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
+    "resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
+    "resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
+    "resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super().__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super().__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False):
+        super().__init__()
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = [block(self.inplanes, planes, stride, downsample)]
+        self.inplanes = planes * block.expansion
+        layers.extend(block(self.inplanes, planes) for _ in range(1, blocks))
+        return nn.Sequential(*layers)
+
+    def forward(self, x, *args):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+def resnet18(pretrained=False, **kwargs):
+    """Constructs a ResNet-18 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls["resnet18"]))
+    return model
+
+
+def resnet34(pretrained=False, **kwargs):
+    """Constructs a ResNet-34 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls["resnet34"]))
+    return model
+
+
+def resnet50(pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls["resnet50"]))
+    return model
+
+
+def resnet101(pretrained=False, **kwargs):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls["resnet101"]))
+    return model
+
+
+def resnet152(pretrained=False, **kwargs):
+    """Constructs a ResNet-152 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls["resnet152"]))
+    return model

AIGVDet/raft_model/raft-things.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcfa4125d6418f4de95d84aec20a3c5f4e205101715a79f193243c186ac9a7e1
+size 21108000

AIGVDet/requirements.txt

@@ -0,0 +1,22 @@
+# conda create -n aigvdet python=3.9
+# pip install torch==2.0.0+cu117 torchvision==0.15.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
+einops
+imageio
+ipympl
+matplotlib
+numpy<2
+opencv-python
+pandas
+scikit-learn
+tensorboard
+tensorboardX
+tqdm
+blobfile>=1.0.5
+natsort
+fastapi==0.116.1
+pydantic==2.11.7
+uvicorn[standard]
+torch==2.0.0+cu117
+torchvision==0.15.1+cu117
+-f https://download.pytorch.org/whl/torch_stable.html
+python-multipart

AIGVDet/run.py

@@ -0,0 +1,214 @@
+import sys
+import argparse
+import os
+import cv2
+import glob
+import numpy as np
+import torch
+from PIL import Image
+import pandas as pd
+import torch
+import torch.nn
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+from tqdm import tqdm
+import time
+
+from .core.raft import RAFT
+from .core.utils import flow_viz
+from .core.utils.utils import InputPadder
+from natsort import natsorted
+from .core.utils1.utils import get_network, str2bool, to_cuda
+from sklearn.metrics import accuracy_score, average_precision_score, roc_auc_score,roc_auc_score
+
+DEVICE = 'cuda'
+# DEVICE = 'cpu'  # Changed to 'cpu'
+device = torch.device(DEVICE)
+
+
+def load_image(imfile):
+    img = np.array(Image.open(imfile)).astype(np.uint8)
+    img = torch.from_numpy(img).permute(2, 0, 1).float()
+    return img[None].to(DEVICE)
+
+
+def viz(img, flo, folder_optical_flow_path, imfile1):
+    img = img[0].permute(1, 2, 0).cpu().numpy()
+    flo = flo[0].permute(1, 2, 0).cpu().numpy()
+    
+    # map flow to RGB image
+    flo = flow_viz.flow_to_image(flo)
+    img_flo = np.concatenate([img, flo], axis=0)
+
+    # extract filename safely (cross-platform)
+    filename = os.path.basename(imfile1).strip()
+    output_path = os.path.join(folder_optical_flow_path, filename)
+
+    print(output_path)
+    cv2.imwrite(output_path, flo)
+
+
+def video_to_frames(video_path, output_folder):
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+    
+    cap = cv2.VideoCapture(video_path)
+    frame_count = 0
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        frame_filename = os.path.join(output_folder, f"frame_{frame_count:05d}.png")
+        cv2.imwrite(frame_filename, frame)
+        frame_count += 1
+    
+    cap.release()
+
+    images = glob.glob(os.path.join(output_folder, '*.png')) + \
+             glob.glob(os.path.join(output_folder, '*.jpg'))
+    images = sorted(images)
+    
+    return images
+
+
+# generate optical flow images
+def OF_gen(args):
+    model = torch.nn.DataParallel(RAFT(args))
+    model.load_state_dict(torch.load(args.model, map_location=torch.device(DEVICE)))
+
+    model = model.module
+    model.to(DEVICE)
+    model.eval()
+
+    if not os.path.exists(args.folder_optical_flow_path):
+        os.makedirs(args.folder_optical_flow_path)
+        print(f'{args.folder_optical_flow_path}')
+
+    with torch.no_grad():
+
+        images = video_to_frames(args.path, args.folder_original_path)
+        images = natsorted(images)
+
+        for imfile1, imfile2 in zip(images[:-1], images[1:]):
+            image1 = load_image(imfile1)
+            image2 = load_image(imfile2)
+
+            padder = InputPadder(image1.shape)
+            image1, image2 = padder.pad(image1, image2)
+
+            flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
+
+            viz(image1, flow_up,args.folder_optical_flow_path,imfile1)
+
+
+def RUN(args=None):
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    default_model_path = os.path.join(script_dir, "raft_model/raft-things.pth")
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model', help="restore checkpoint",default=default_model_path)
+    parser.add_argument('--path', help="dataset for evaluation",default="video/000000.mp4")
+    parser.add_argument('--folder_original_path', help="dataset for evaluation_frames",default="frame/000000")
+    parser.add_argument('--small', action='store_true', help='use small model')
+    parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
+    parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation')
+    parser.add_argument('--folder_optical_flow_path',help="the results to save",default="optical_result/000000")
+    parser.add_argument(
+        "-mop",
+        "--model_optical_flow_path",
+        type=str,
+        default="checkpoints/optical.pth",
+    )
+    parser.add_argument(
+        "-mor",
+        "--model_original_path",
+        type=str,
+        default="checkpoints/original.pth",
+    )
+    parser.add_argument(
+        "-t",
+        "--threshold",
+        type=float,
+        default=0.5,
+    )
+    parser.add_argument("--use_cpu", action="store_true", help="uses gpu by default, turn on to use cpu")
+    parser.add_argument("--arch", type=str, default="resnet50")
+    parser.add_argument("--aug_norm", type=str2bool, default=True)
+    args = parser.parse_args(args)
+    start_time = time.perf_counter()
+    OF_gen(args)
+    elapsed = time.perf_counter() - start_time
+    print(f"⏱️ [OF_gen] Service call took {elapsed:.2f} seconds")
+    # Load models
+    model_op = get_network(args.arch)
+    state_dict = torch.load(args.model_optical_flow_path, map_location='cpu')
+    model_op.load_state_dict(state_dict["model"] if "model" in state_dict else state_dict)
+    model_op.eval().to(device)
+
+    model_or = get_network(args.arch)
+    state_dict = torch.load(args.model_original_path, map_location='cpu')
+    model_or.load_state_dict(state_dict["model"] if "model" in state_dict else state_dict)
+    model_or.eval().to(device)
+
+    # Transform
+    trans = transforms.Compose([
+        transforms.CenterCrop((448, 448)),
+        transforms.ToTensor(),
+    ])
+
+    # Process original frames
+    original_file_list = sorted(
+        glob.glob(os.path.join(args.folder_original_path, "*.jpg")) +
+        glob.glob(os.path.join(args.folder_original_path, "*.png")) +
+        glob.glob(os.path.join(args.folder_original_path, "*.JPEG"))
+    )
+    original_prob_sum = 0
+    for img_path in tqdm(original_file_list, desc="Original", dynamic_ncols=True, disable=len(original_file_list) <= 1):
+        img = Image.open(img_path).convert("RGB")
+        img = trans(img)
+        if args.aug_norm:
+            img = TF.normalize(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        in_tens = img.unsqueeze(0).to(device)
+        with torch.no_grad():
+            prob = model_or(in_tens).sigmoid().item()
+        original_prob_sum += prob
+
+    original_prob = original_prob_sum / len(original_file_list)
+    print(f"Original prob: {original_prob:.4f}")
+
+    # Process optical flow frames
+    optical_file_list = sorted(
+        glob.glob(os.path.join(args.folder_optical_flow_path, "*.jpg")) +
+        glob.glob(os.path.join(args.folder_optical_flow_path, "*.png")) +
+        glob.glob(os.path.join(args.folder_optical_flow_path, "*.JPEG"))
+    )
+    optical_prob_sum = 0
+    for img_path in tqdm(optical_file_list, desc="Optical Flow", dynamic_ncols=True, disable=len(optical_file_list) <= 1):
+        img = Image.open(img_path).convert("RGB")
+        img = trans(img)
+        if args.aug_norm:
+            img = TF.normalize(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        in_tens = img.unsqueeze(0).to(device)
+        with torch.no_grad():
+            prob = model_op(in_tens).sigmoid().item()
+        optical_prob_sum += prob
+
+    optical_prob = optical_prob_sum / len(optical_file_list)
+    print(f"Optical prob: {optical_prob:.4f}")
+
+    final_prob = (original_prob + optical_prob) / 2
+    print(f"predict: {final_prob}")
+    
+    real_score = 1 - final_prob
+    fake_score = final_prob
+    print(f"Confidence scores - Real: {real_score:.4f}, Fake: {fake_score:.4f}")
+
+    return {
+        "original_prob": original_prob,
+        "optical_prob": optical_prob,
+        "final_predict": final_prob,
+        "real_score": real_score,
+        "fake_score": fake_score
+    }

AIGVDet/run.sh

@@ -0,0 +1 @@
+python demo.py --path "demo_video/real/094.mp4" --folder_original_path "frame/000000" --folder_optical_flow_path "optical_result/000000" -mop "checkpoints/optical.pth" -mor "checkpoints/original.pth"

AIGVDet/test.py

@@ -0,0 +1,261 @@
+import argparse
+import glob
+import os
+import pandas as pd
+
+import torch
+import torch.nn
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+from PIL import Image
+from tqdm import tqdm
+
+from core.utils1.utils import get_network, str2bool, to_cuda
+from sklearn.metrics import accuracy_score, average_precision_score, roc_auc_score,roc_auc_score
+
+if __name__=="__main__":
+
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "-fop", "--folder_optical_flow_path", default="data/test/T2V/videocraft", type=str, help="path to optical flow imagefile folder"
+    )
+    parser.add_argument(
+        "-for", "--folder_original_path", default="data/test/original/T2V/videocraft", type=str, help="path to RGB image file folder"
+    )
+    parser.add_argument(
+        "-mop",
+        "--model_optical_flow_path",
+        type=str,
+        default="checkpoints/optical.pth",
+    )
+    parser.add_argument(
+        "-mor",
+        "--model_original_path",
+        type=str,
+        default="checkpoints/original.pth",
+    )
+
+    parser.add_argument(
+        "-t",
+        "--threshold",
+        type=float,
+        default=0.5,
+    )
+    
+    parser.add_argument(
+        "-e",
+        "--excel_path",
+        type=str,
+        help="path to excel of frames",
+        default="data/results/moonvalley_wang.csv",
+    )
+    
+    parser.add_argument(
+        "-ef",
+        "--excel_frame_path",
+        type=str,
+        help="path to excel of frame detection result",
+        default="data/results/frame/moonvalley_wang.csv",
+    )
+    
+    
+    
+    
+    parser.add_argument("--use_cpu", action="store_true", help="uses gpu by default, turn on to use cpu")
+    parser.add_argument("--arch", type=str, default="resnet50")
+    parser.add_argument("--aug_norm", type=str2bool, default=True)
+
+    args = parser.parse_args()
+    subfolder_count = 0
+
+    model_op = get_network(args.arch)
+    state_dict = torch.load(args.model_optical_flow_path, map_location="cpu")
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+    model_op.load_state_dict(state_dict)
+    model_op.eval()
+    if not args.use_cpu:
+        model_op.cuda()
+        
+        
+    model_or = get_network(args.arch)
+    state_dict = torch.load(args.model_original_path, map_location="cpu")
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+    model_or.load_state_dict(state_dict)
+    model_or.eval()
+    if not args.use_cpu:
+        model_or.cuda()
+    
+    
+    trans = transforms.Compose(
+        (
+            transforms.CenterCrop((448,448)),
+            transforms.ToTensor(),
+        )
+    )
+
+    print("*" * 50)
+
+    flag=0
+    p=0
+    n=0
+    tp=0
+    tn=0
+    y_true=[]
+    y_pred=[]
+
+    # create an empty DataFrame
+    df = pd.DataFrame(columns=['name', 'pro','flag','optical_pro','original_pro'])
+    df1 = pd.DataFrame(columns=['original_path', 'original_pro','optical_path','optical_pro','flag'])
+    index1=0
+    
+    # Traverse through subfolders in a large folder.
+    for subfolder_name in ["0_real", "1_fake"]:
+        optical_subfolder_path = os.path.join(args.folder_optical_flow_path, subfolder_name)
+        original_subfolder_path = os.path.join(args.folder_original_path, subfolder_name)
+
+        if subfolder_name=="0_real":
+            flag=0
+        else:
+            flag=1
+            
+        if os.path.isdir(optical_subfolder_path):
+            pass
+        else:
+            print("Subfolder does not exist.", optical_subfolder_path)
+        
+        # Check if the subfolder path exists.
+        if os.path.isdir(original_subfolder_path):
+            print("test subfolder:", subfolder_name)
+
+            # Traverse through sub-subfolders within a subfolder.
+            for subsubfolder_name in os.listdir(original_subfolder_path):
+                original_subsubfolder_path = os.path.join(original_subfolder_path, subsubfolder_name)
+                optical_subsubfolder_path = os.path.join(optical_subfolder_path, subsubfolder_name)
+                if os.path.isdir(optical_subsubfolder_path):
+                    pass
+                else:
+                    print("Sub-subfolder does not exist.",optical_subsubfolder_path)
+                    
+                if os.path.isdir(original_subsubfolder_path):
+                    print("test subsubfolder:", subsubfolder_name)
+                    
+                    #Detect original
+                    original_file_list = sorted(glob.glob(os.path.join(original_subsubfolder_path, "*.jpg")) + glob.glob(os.path.join(original_subsubfolder_path, "*.png"))+glob.glob(os.path.join(original_subsubfolder_path, "*.JPEG")))
+
+                    original_prob_sum=0
+                    for img_path in tqdm(original_file_list, dynamic_ncols=True, disable=len(original_file_list) <= 1):
+                        
+                        img = Image.open(img_path).convert("RGB")
+                        img = trans(img)
+                        if args.aug_norm:
+                            img = TF.normalize(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+                        in_tens = img.unsqueeze(0)
+                        if not args.use_cpu:
+                            in_tens = in_tens.cuda()
+                        
+                        with torch.no_grad():
+                            prob = model_or(in_tens).sigmoid().item()
+                            original_prob_sum+=prob
+                            
+                        df1 = df1.append({'original_path': img_path, 'original_pro': prob , 'flag':flag}, ignore_index=True)
+                        
+                        
+                    original_predict=original_prob_sum/len(original_file_list)
+                    print("original prob",original_predict)
+                    
+                    #Detect optical flow
+                    optical_file_list = sorted(glob.glob(os.path.join(optical_subsubfolder_path, "*.jpg")) + glob.glob(os.path.join(optical_subsubfolder_path, "*.png"))+glob.glob(os.path.join(optical_subsubfolder_path, "*.JPEG")))
+                    optical_prob_sum=0
+                    for img_path in tqdm(optical_file_list, dynamic_ncols=True, disable=len(original_file_list) <= 1):
+                        
+                        img = Image.open(img_path).convert("RGB")
+                        img = trans(img)
+                        if args.aug_norm:
+                            img = TF.normalize(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+                        in_tens = img.unsqueeze(0)
+                        if not args.use_cpu:
+                            in_tens = in_tens.cuda()
+
+                        with torch.no_grad():
+                            prob = model_op(in_tens).sigmoid().item()
+                            optical_prob_sum+=prob
+                        
+                        df1.loc[index1, 'optical_path'] = img_path
+                        df1.loc[index1, 'optical_pro'] = prob
+                        index1=index1+1
+                    index1=index1+1
+                    
+                    optical_predict=optical_prob_sum/len(optical_file_list)
+                    print("optical prob",optical_predict)
+                    
+                    predict=original_predict*0.5+optical_predict*0.5
+                    print(f"flag:{flag} predict:{predict}")
+                    # y_true.append((float)(flag))
+                    y_true.append((flag))
+                    y_pred.append(predict)
+                    if flag==0:
+                        n+=1
+                        if predict<args.threshold:
+                            tn+=1
+                    else:
+                        p+=1
+                        if predict>=args.threshold:
+                            tp+=1
+                    df = df.append({'name': subsubfolder_name, 'pro': predict , 'flag':flag ,'optical_pro':optical_predict,'original_pro':original_predict}, ignore_index=True)
+        else:
+            print("Subfolder does not exist:", original_subfolder_path)
+    # r_acc = accuracy_score(y_true[y_true == 0], y_pred[y_true == 0] > args.threshold)
+    # f_acc = accuracy_score(y_true[y_true == 1], y_pred[y_true == 1] > args.threshold)
+    # acc = accuracy_score(y_true, y_pred > args.threshold)
+    
+    ap = average_precision_score(y_true, y_pred)
+    auc=roc_auc_score(y_true,y_pred)
+    # print(f"r_acc:{r_acc}")
+    print(f"tnr:{tn/n}")
+    # print(f"f_acc:{f_acc}")
+    print(f"tpr:{tp/p}")
+    print(f"acc:{(tp+tn)/(p+n)}")
+    # print(f"acc:{acc}")
+    print(f"ap:{ap}")
+    print(f"auc:{auc}")
+    print(f"p:{p}")
+    print(f"n:{n}")
+    print(f"tp:{tp}")
+    print(f"tn:{tn}")
+
+    # Write the DataFrame to a csv file.
+    csv_filename = args.excel_path
+    csv_folder = os.path.dirname(csv_filename) 
+    if not os.path.exists(csv_folder):
+        os.makedirs(csv_folder)
+
+
+    if not os.path.exists(csv_filename):
+        df.to_csv(csv_filename, index=False)
+    else:
+        df.to_csv(csv_filename, mode='a', header=False, index=False)
+    print(f"Results have been saved to {csv_filename}")
+        
+    # Write the prediction probabilities of the frame to a CSV file.
+    csv_filename1 = args.excel_frame_path
+    csv_folder1 = os.path.dirname(csv_filename1) 
+    if not os.path.exists(csv_folder1):
+        os.makedirs(csv_folder1)
+
+    if not os.path.exists(csv_filename1):
+        df1.to_csv(csv_filename1, index=False)
+    else:
+        df1.to_csv(csv_filename1, mode='a', header=False, index=False)
+    
+    # if not os.path.exists(excel_filename):
+    #     with pd.ExcelWriter(excel_filename, engine='xlsxwriter') as writer:
+    #         df.to_excel(writer, sheet_name='Sheet1', index=False)
+    # else:
+    #     with pd.ExcelWriter(excel_filename, mode='a', engine='openpyxl') as writer:
+    #         df.to_excel(writer, sheet_name='Sheet1', index=False, startrow=0, header=False)
+    print(f"Results have been saved to {csv_filename1}")
+    
+
+

AIGVDet/test.sh

@@ -0,0 +1 @@
+python test.py -fop "data/test/T2V/hotshot" -mop "checkpoints/optical_aug.pth" -for "data/test/original/T2V/hotshot" -mor "checkpoints/original_aug.pth" -e "data/results/T2V/hotshot.csv" -ef "data/results/frame/T2V/hotshot.csv" -t 0.5

AIGVDet/train.py

@@ -0,0 +1,87 @@
+from core.utils1.config import cfg  # isort: split
+
+import os
+import time
+
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+from core.utils1.datasets import create_dataloader
+from core.utils1.earlystop import EarlyStopping
+from core.utils1.eval import get_val_cfg, validate
+from core.utils1.trainer import Trainer
+from core.utils1.utils import Logger
+
+import ssl
+ssl._create_default_https_context = ssl._create_unverified_context
+
+
+if __name__ == "__main__":
+    val_cfg = get_val_cfg(cfg, split="val", copy=True)
+    cfg.dataset_root = os.path.join(cfg.dataset_root, "train")
+    data_loader = create_dataloader(cfg)
+    dataset_size = len(data_loader)
+
+    log = Logger()
+    log.open(cfg.logs_path, mode="a")
+    log.write("Num of training images = %d\n" % (dataset_size * cfg.batch_size))
+    log.write("Config:\n" + str(cfg.to_dict()) + "\n")
+
+    train_writer = SummaryWriter(os.path.join(cfg.exp_dir, "train"))
+    val_writer = SummaryWriter(os.path.join(cfg.exp_dir, "val"))
+
+    trainer = Trainer(cfg)
+    early_stopping = EarlyStopping(patience=cfg.earlystop_epoch, delta=-0.001, verbose=True)
+    for epoch in range(cfg.nepoch):
+        epoch_start_time = time.time()
+        iter_data_time = time.time()
+        epoch_iter = 0
+
+        for data in tqdm(data_loader, dynamic_ncols=True):
+            trainer.total_steps += 1
+            epoch_iter += cfg.batch_size
+
+            trainer.set_input(data)
+            trainer.optimize_parameters()
+
+            # if trainer.total_steps % cfg.loss_freq == 0:
+            #     log.write(f"Train loss: {trainer.loss} at step: {trainer.total_steps}\n")
+            train_writer.add_scalar("loss", trainer.loss, trainer.total_steps)
+
+            if trainer.total_steps % cfg.save_latest_freq == 0:
+                log.write(
+                    "saving the latest model %s (epoch %d, model.total_steps %d)\n"
+                    % (cfg.exp_name, epoch, trainer.total_steps)
+                )
+                trainer.save_networks("latest")
+
+        if epoch % cfg.save_epoch_freq == 0:
+            log.write("saving the model at the end of epoch %d, iters %d\n" % (epoch, trainer.total_steps))
+            trainer.save_networks("latest")
+            trainer.save_networks(epoch)
+
+        # Validation
+        trainer.eval()
+        val_results = validate(trainer.model, val_cfg)
+        val_writer.add_scalar("AP", val_results["AP"], trainer.total_steps)
+        val_writer.add_scalar("ACC", val_results["ACC"], trainer.total_steps)
+        # add
+        val_writer.add_scalar("AUC", val_results["AUC"], trainer.total_steps)
+        val_writer.add_scalar("TPR", val_results["TPR"], trainer.total_steps)
+        val_writer.add_scalar("TNR", val_results["TNR"], trainer.total_steps)
+        
+        log.write(f"(Val @ epoch {epoch}) AP: {val_results['AP']}; ACC: {val_results['ACC']}\n")
+
+        if cfg.earlystop:
+            early_stopping(val_results["ACC"], trainer)
+            if early_stopping.early_stop:
+                if trainer.adjust_learning_rate():
+                    log.write("Learning rate dropped by 10, continue training...\n")
+                    early_stopping = EarlyStopping(patience=cfg.earlystop_epoch, delta=-0.002, verbose=True)
+                else:
+                    log.write("Early stopping.\n")
+                    break
+        if cfg.warmup:
+            # print(trainer.scheduler.get_lr()[0])
+            trainer.scheduler.step()
+        trainer.train()

AIGVDet/train.sh

@@ -0,0 +1,4 @@
+EXP_NAME="moonvalley_vos2_crop"
+DATASETS="moonvalley_vos2_crop"
+DATASETS_TEST="moonvalley_vos2_crop"
+python train.py --gpus 0 --exp_name $EXP_NAME datasets $DATASETS datasets_test $DATASETS_TEST 

api_server.py

@@ -0,0 +1,281 @@
+import asyncio
+import os
+import shutil
+import json
+import uuid
+import time
+import threading
+from typing import List, Dict, Any, Optional
+from concurrent.futures import ThreadPoolExecutor
+
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+# Giả lập import các thư viện của bạn
+from miragenews import run_multimodal_to_json
+from AIGVDet import run_video_to_json
+
+UPLOAD_DIR = "temp_uploads"
+MAX_WORKERS = 4  
+
+app = FastAPI(
+    title="Multimedia Analysis API (Polling Mode)",
+    description="API phân tích đa phương tiện sử dụng cơ chế Polling để tránh Timeout.",
+    version="2.0.0",
+)
+
+jobs: Dict[str, Dict] = {}
+jobs_lock = threading.Lock()
+
+class AnalysisResult(BaseModel):
+    image_analysis_results: Optional[List[Any]] = None
+    video_analysis_result: Optional[Dict[str, Any]] = None
+
+class JobStatus(BaseModel):
+    job_id: str
+    status: str  
+    message: Optional[str] = None
+    result: Optional[AnalysisResult] = None
+    created_at: float
+    updated_at: float
+
+def _create_job() -> str:
+    job_id = uuid.uuid4().hex
+    now = time.time()
+    with jobs_lock:
+        jobs[job_id] = {
+            "job_id": job_id,
+            "status": "queued",
+            "message": "Đang chờ xử lý...",
+            "result": None,
+            "created_at": now,
+            "updated_at": now
+        }
+    return job_id
+
+def _update_job(job_id: str, **kwargs):
+    with jobs_lock:
+        if job_id in jobs:
+            jobs[job_id].update(kwargs)
+            jobs[job_id]["updated_at"] = time.time()
+
+def _get_job(job_id: str) -> Optional[Dict]:
+    with jobs_lock:
+        return jobs.get(job_id)
+
+async def run_analysis_logic(
+    image_paths: Optional[List[str]] = None,
+    video_path: Optional[str] = None,
+    text: str = "",
+) -> Dict[str, Any]:
+    
+    if not image_paths and not video_path:
+        raise ValueError("Cần cung cấp ít nhất một trong hai: image_paths hoặc video_path.")
+
+    tasks = []
+    
+    if image_paths:
+        image_task = asyncio.create_task(
+            run_multimodal_to_json(image_paths=image_paths, text=text, output_json_path=None)
+        )
+        tasks.append(image_task)
+
+    if video_path:
+        video_task = asyncio.to_thread(
+            run_video_to_json, video_path=video_path, output_json_path=None
+        )
+        tasks.append(video_task)
+
+    task_results = await asyncio.gather(*tasks)
+
+    final_result = {"image_analysis_results": [], "video_analysis_result": {}}
+    image_analysis_results = []
+    video_result_index = -1
+
+    current_idx = 0
+    if image_paths:
+        image_analysis_results = task_results[current_idx]
+        current_idx += 1
+    
+    if video_path:
+        video_result_index = current_idx
+
+    final_result["image_analysis_results"] = image_analysis_results
+
+    if video_result_index != -1:
+        raw_video_result = task_results[video_result_index]
+        if raw_video_result:
+            video_id_key = list(raw_video_result.keys())[0]
+            video_data = raw_video_result[video_id_key]
+            
+            avg_authentic = video_data.get("authentic_confidence_score", 0)
+            avg_synthetic = video_data.get("synthetic_confidence_score", 0)
+
+            if avg_authentic > avg_synthetic and avg_authentic > 0.5:
+                authenticity_assessment = "REAL (Authentic)"
+                verification_tools = "Deepfake Detector"
+                synthetic_type = "N/A"
+                other_artifacts = "Our algorithms conducted a thorough analysis of the video's motion patterns, lighting consistency, and object interactions. We observed fluid, natural movements and consistent physics that align with real-world recordings. No discernible artifacts, such as pixel distortion, unnatural blurring, or shadow inconsistencies, were detected that would indicate digital manipulation or AI-driven synthesis."
+            elif avg_authentic > avg_synthetic and avg_authentic <= 0.5:
+                authenticity_assessment = "Potentially Synthetic"
+                verification_tools = "Deepfake Detector"
+                synthetic_type = "Potentially AI-generated"
+                other_artifacts = "Our analysis has identified subtle anomalies within the video frames, particularly in areas of complex texture and inconsistent lighting across different objects. While these discrepancies are not significant enough to definitively classify the video as synthetic, they do suggest a possibility of digital alteration or partial AI generation. Further examination may be required for a conclusive determination."
+            else:
+                authenticity_assessment = "NOT REAL (Fake, Manipulated, or AI)"
+                verification_tools = "Deepfake Detector"
+                synthetic_type = "AI-generated"
+                other_artifacts = "Our deep analysis detected multiple, significant artifacts commonly associated with synthetic or manipulated media. These include, but are not limited to, unnatural facial expressions and eye movements, inconsistent or floating shadows, logical impossibilities in object interaction, and high-frequency digital noise characteristic of generative models. These factors strongly indicate that the video is not authentic."
+
+            final_result["video_analysis_result"] = {
+                "filename": video_data.get("video_name", ""),
+                "result": {
+                    "authenticity_assessment": authenticity_assessment,
+                    "verification_tools_methods": verification_tools,
+                    "synthetic_type": synthetic_type,
+                    "other_artifacts": other_artifacts,
+                },
+            }
+
+    if not final_result.get("image_analysis_results"):
+        final_result.pop("image_analysis_results", None)
+    if not final_result.get("video_analysis_result"):
+        final_result.pop("video_analysis_result", None)
+
+    return final_result
+
+async def process_job_background(
+    job_id: str, 
+    temp_dir: str, 
+    image_paths: List[str], 
+    video_path: str, 
+    text: str
+):
+    """Hàm chạy ngầm thực hiện phân tích"""
+    _update_job(job_id, status="running", message="Đang phân tích...")
+    
+    try:
+        result_data = await run_analysis_logic(
+            image_paths=image_paths if image_paths else None,
+            video_path=video_path if video_path else None,
+            text=text
+        )
+        
+        _update_job(job_id, status="succeeded", result=result_data, message="Hoàn tất")
+        
+    except Exception as e:
+        print(f"Error processing job {job_id}: {e}")
+        _update_job(job_id, status="failed", message=str(e))
+        
+    finally:
+        try:
+            if os.path.exists(temp_dir):
+                shutil.rmtree(temp_dir)
+                print(f"Deleted temp dir: {temp_dir}")
+        except Exception as cleanup_error:
+            print(f"Cleanup error for {job_id}: {cleanup_error}")
+
+
+@app.on_event("startup")
+def startup_event():
+    os.makedirs(UPLOAD_DIR, exist_ok=True)
+
+@app.get("/analyze/{job_id}", response_model=JobStatus)
+async def get_job_status(job_id: str):
+    """Client gọi API này định kỳ để kiểm tra kết quả"""
+    job = _get_job(job_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
+    return job
+
+@app.post("/analyze/image/", response_model=JobStatus)
+async def analyze_image_endpoint(
+    background_tasks: BackgroundTasks,
+    images: List[UploadFile] = File(...),
+    text: Optional[str] = Form(""),
+):
+    job_id = _create_job()
+    
+    job_dir = os.path.join(UPLOAD_DIR, job_id)
+    os.makedirs(job_dir, exist_ok=True)
+    
+    saved_image_paths = []
+    try:
+        for img in images:
+            file_path = os.path.join(job_dir, img.filename)
+            with open(file_path, "wb") as buffer:
+                shutil.copyfileobj(img.file, buffer)
+            saved_image_paths.append(file_path)
+    except Exception as e:
+        _update_job(job_id, status="failed", message=f"Lỗi upload: {e}")
+        return _get_job(job_id)
+
+    background_tasks.add_task(
+        process_job_background, 
+        job_id, job_dir, saved_image_paths, None, text
+    )
+
+    return _get_job(job_id)
+
+@app.post("/analyze/video/", response_model=JobStatus)
+async def analyze_video_endpoint(
+    background_tasks: BackgroundTasks,
+    video: UploadFile = File(...),
+):
+    job_id = _create_job()
+    job_dir = os.path.join(UPLOAD_DIR, job_id)
+    os.makedirs(job_dir, exist_ok=True)
+    
+    saved_video_path = os.path.join(job_dir, video.filename)
+    try:
+        with open(saved_video_path, "wb") as buffer:
+            shutil.copyfileobj(video.file, buffer)
+    except Exception as e:
+         _update_job(job_id, status="failed", message=f"Lỗi upload: {e}")
+         return _get_job(job_id)
+
+    background_tasks.add_task(
+        process_job_background, 
+        job_id, job_dir, [], saved_video_path, ""
+    )
+
+    return _get_job(job_id)
+
+@app.post("/analyze/multimodal/", response_model=JobStatus)
+async def analyze_multimodal_endpoint(
+    background_tasks: BackgroundTasks,
+    images: List[UploadFile] = File(...),
+    video: UploadFile = File(...),
+    text: Optional[str] = Form(""),
+):
+    job_id = _create_job()
+    job_dir = os.path.join(UPLOAD_DIR, job_id)
+    os.makedirs(job_dir, exist_ok=True)
+    
+    saved_image_paths = []
+    saved_video_path = None
+    
+    try:
+        # Save Images
+        for img in images:
+            file_path = os.path.join(job_dir, img.filename)
+            with open(file_path, "wb") as buffer:
+                shutil.copyfileobj(img.file, buffer)
+            saved_image_paths.append(file_path)
+        
+        # Save Video
+        saved_video_path = os.path.join(job_dir, video.filename)
+        with open(saved_video_path, "wb") as buffer:
+            shutil.copyfileobj(video.file, buffer)
+            
+    except Exception as e:
+        _update_job(job_id, status="failed", message=f"Lỗi upload: {e}")
+        return _get_job(job_id)
+
+    background_tasks.add_task(
+        process_job_background, 
+        job_id, job_dir, saved_image_paths, saved_video_path, text
+    )
+
+    return _get_job(job_id)

checkpoints/image/best-mirage-img.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:755e26bfe97830e03e4c475596ade7168e1df18c57d89162d685b1148ae9f5a8
+size 2375

checkpoints/image/cbm-encoder.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a90c02450ae4a105414c6701bf571d1a498571083a725d916119d792cbdcf2f
+size 1917947

checkpoints/image/cbm-predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2eba7ff0cd8db1c69961478551b2a40e47687f281d7e853f95a6650908e8c3df
+size 2387